Update app.py

app.py

@@ -1,27 +1,63 @@
-# app.py
 import streamlit as st
 import soundfile as sf
 import torch
 from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-from train_pipeline import train_model
-
-st.title("Telugu ASR - XLS-R Training + Inference")
-
-if st.button("Start Training"):
-    st.write("🚀 Training started…")
-    train_model()
-    st.success("Training complete!")
-
-uploaded = st.file_uploader("Upload WAV audio", type=["wav"])
-if uploaded:
-    audio, sr = sf.read(uploaded)
-    processor = Wav2Vec2Processor.from_pretrained("./model")
-    model = Wav2Vec2ForCTC.from_pretrained("./model").to("cpu")
-
-    inputs = processor(audio, sampling_rate=sr, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(inputs.input_values).logits
-        
-    pred_ids = torch.argmax(logits, dim=-1)
-    st.subheader("Transcription")
-    st.write(processor.batch_decode(pred_ids)[0])
+import os
+import subprocess
+
+st.set_page_config(page_title="Telugu ASR Training", layout="wide")
+
+st.title("Telugu ASR - Train XLS-R + Run Inference")
+
+# ============================================================
+# SECTION 1 — TRAIN BUTTON
+# ============================================================
+
+st.header("Train the Model")
+
+st.write("""
+Click the button below to start training the Wav2Vec2-XLSR model on Telugu dataset.
+This will run train_pipeline.py inside this Space.
+""")
+
+if st.button("🚀 Start Training"):
+    st.info("Training started... this may take several minutes to hours depending on GPU.")
+
+    # run training script as a subprocess so Streamlit initializes properly
+    process = subprocess.Popen(
+        ["python3", "train_pipeline.py"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True
+    )
+
+    st.success("Training script launched. Check the Space logs for progress.")
+
+# ============================================================
+# SECTION 2 — INFERENCE
+# ============================================================
+
+st.header("Inference (Upload Audio)")
+
+if not os.path.exists("./model"):
+    st.warning("No model found. Train it first.")
+else:
+    uploaded_audio = st.file_uploader("Upload WAV File", type=["wav"])
+
+    if uploaded_audio is not None:
+        audio, sr = sf.read(uploaded_audio)
+
+        # Load trained model
+        processor = Wav2Vec2Processor.from_pretrained("./model")
+        model = Wav2Vec2ForCTC.from_pretrained("./model")
+
+        inputs = processor(audio, sampling_rate=sr, return_tensors="pt", padding=True)
+
+        with torch.no_grad():
+            logits = model(inputs.input_values).logits
+
+        pred_ids = torch.argmax(logits, dim=-1)
+        transcription = processor.batch_decode(pred_ids)[0]
+
+        st.subheader("Transcription")
+        st.write(transcription)