Update app.py

app.py

@@ -6,20 +6,14 @@ import torch
 import librosa
 import soundfile as sf
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForCTC, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoProcessor, AutoModelForCTC
 
-CHUNK_SEC      = 20
-OVERLAP_SEC    = 1.0
-MT_BATCH_SIZE  = 8
-MODEL_ID       = os.environ.get("MODEL_ID")
-MT_MODEL_ID    = os.environ.get("MT_MODEL_ID")
-SRC_LANG       = "nya_Latn"
-TGT_LANG       = "eng_Latn"
+CHUNK_SEC   = 20
+OVERLAP_SEC = 1.0
+MODEL_ID    = os.environ.get("MODEL_ID")
 
-processor  = AutoProcessor.from_pretrained(MODEL_ID)
-model      = AutoModelForCTC.from_pretrained(MODEL_ID).eval()
-mt_tok     = AutoTokenizer.from_pretrained(MT_MODEL_ID, src_lang=SRC_LANG)
-mt_model   = AutoModelForSeq2SeqLM.from_pretrained(MT_MODEL_ID).eval()
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model     = AutoModelForCTC.from_pretrained(MODEL_ID).eval()
 
 
 def convert_to_wav(input_path):
@@ -30,9 +24,8 @@ def convert_to_wav(input_path):
 
 
 @spaces.GPU
-def transcribe_and_translate(audio_path, progress=gr.Progress()):
+def transcribe(audio_path, progress=gr.Progress()):
     model.to("cuda")
-    mt_model.to("cuda")
 
     wav       = convert_to_wav(audio_path)
     audio, sr = librosa.load(wav, sr=16000, mono=True)
@@ -66,71 +59,26 @@ def transcribe_and_translate(audio_path, progress=gr.Progress()):
 
     transcript = " ".join(parts).strip()
 
-    # Batched translation
-    progress(1.0, desc="Translating...")
-    sentences = [s.strip() for s in transcript.replace("\n", ". ").split(".") if s.strip()]
-    translated_parts = []
-    for batch_start in range(0, len(sentences), MT_BATCH_SIZE):
-        batch = sentences[batch_start : batch_start + MT_BATCH_SIZE]
-        inp   = mt_tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        inp   = {k: v.to("cuda") for k, v in inp.items()}
-        with torch.no_grad():
-            out = mt_model.generate(
-                **inp,
-                forced_bos_token_id=mt_tok.convert_tokens_to_ids(TGT_LANG))
-        translated_parts.extend(mt_tok.batch_decode(out, skip_special_tokens=True))
-
-    translation = " ".join(translated_parts)
-
     txt = tempfile.mktemp(suffix=".txt")
     with open(txt, "w", encoding="utf-8") as f:
-        f.write("=== TRANSCRIPT ===\n")
-        f.write(transcript + "\n\n")
-        f.write("=== TRANSLATION ===\n")
-        f.write(translation + "\n")
-
-    return transcript, translation, gr.update(value=txt, visible=True)
-
-
-@spaces.GPU
-def translate_text(text):
-    mt_model.to("cuda")
-
-    sentences = [s.strip() for s in text.replace("\n", ". ").split(".") if s.strip()]
-    translated_parts = []
-    for batch_start in range(0, len(sentences), MT_BATCH_SIZE):
-        batch = sentences[batch_start : batch_start + MT_BATCH_SIZE]
-        inp   = mt_tok(batch, return_tensors="pt", padding=True, truncation=True, max_length=512)
-        inp   = {k: v.to("cuda") for k, v in inp.items()}
-        with torch.no_grad():
-            out = mt_model.generate(
-                **inp,
-                forced_bos_token_id=mt_tok.convert_tokens_to_ids(TGT_LANG))
-        translated_parts.extend(mt_tok.batch_decode(out, skip_special_tokens=True))
+        f.write(transcript + "\n")
 
-    return " ".join(translated_parts)
+    return transcript, gr.update(value=txt, visible=True)
 
 
 with gr.Blocks(title="Chichewa ASR") as demo:
-    gr.Markdown("## Chichewa Speech Transcription and Translation")
-
-    with gr.Tabs():
-        with gr.Tab("Transcribe"):
-            audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath")
-            transcribe_btn = gr.Button("Process")
-            transcript       = gr.Textbox(label="Speech Transcription", lines=3, interactive=False)
-            transcript_trans = gr.Textbox(label="English Translation", lines=3, interactive=False)
-
-        with gr.Tab("Translate"):
-            text_input    = gr.Textbox(label="Chichewa Text", lines=6)
-            translate_btn = gr.Button("Translate")
-            text_output   = gr.Textbox(label="English Translation", lines=6, interactive=False)
-
-    def run_transcribe(audio):
-        t, tr, _ = transcribe_and_translate(audio)
-        return t, tr
-
-    transcribe_btn.click(fn=run_transcribe, inputs=audio_input, outputs=[transcript, transcript_trans], show_progress="full")
-    translate_btn.click(fn=translate_text, inputs=text_input, outputs=text_output, show_progress="full")
+    gr.Markdown("## Chichewa Speech Transcription")
+
+    audio_input    = gr.Audio(sources=["upload", "microphone"], type="filepath")
+    transcribe_btn = gr.Button("Transcribe")
+    transcript_out = gr.Textbox(label="Transcription", lines=6, interactive=False)
+    download_out   = gr.File(label="Download transcript", visible=False)
+
+    transcribe_btn.click(
+        fn=transcribe,
+        inputs=audio_input,
+        outputs=[transcript_out, download_out],
+        show_progress="full",
+    )
 
 demo.launch(debug=True)