Updated app.py

app.py

@@ -4,7 +4,6 @@ import numpy as np
 import librosa
 import os
 import ctranslate2
-
 from transformers import (
     AutoProcessor,
     AutoModelForSpeechSeq2Seq,
@@ -35,36 +34,40 @@ asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(
 ).to(device)
 
 asr_model.eval()
-
 print("ASR Loaded")
 
 # ====================================
-# LOAD TOKENIZER
+# LOAD MT TOKENIZER
 # ====================================
-tokenizer = AutoTokenizer.from_pretrained(MT_MODEL)
+print("Loading MT tokenizer...")
+mt_tokenizer = AutoTokenizer.from_pretrained(MT_MODEL)
+print("MT tokenizer loaded")
 
 # ====================================
-# CT2 CONVERT
+# CT2 CONVERT & LOAD
 # ====================================
 if not os.path.exists(CT2_DIR):
+    print("Converting MT model to CTranslate2 format...")
     os.system(
         f"ct2-transformers-converter "
         f"--model {MT_MODEL} "
         f"--output_dir {CT2_DIR} "
         f"--quantization int8"
     )
+    print("Conversion done")
 
+print("Loading CTranslate2 translator...")
 translator = ctranslate2.Translator(
     CT2_DIR,
     device=device,
     compute_type="int8"
 )
+print("Translator loaded")
 
 # ====================================
-# AUDIO
+# AUDIO UTILS
 # ====================================
 def fix_audio(audio):
-
     sr, wav = audio
 
     if len(wav.shape) > 1:
@@ -81,12 +84,10 @@ def fix_audio(audio):
 
     return wav
 
-
 # ====================================
 # ASR
 # ====================================
 def transcribe(audio):
-
     if audio is None:
         return ""
 
@@ -114,51 +115,47 @@ def transcribe(audio):
 
     return text
 
-
 # ====================================
 # MT
 # ====================================
 def translate(text):
-
     if not text:
         return ""
 
-    tokenizer.src_lang = "iboLatn"
+    # Prepend source tag directly to text, matching the working API
+    input_text = f"ibo_Latn {text}"
 
-    ids = tokenizer.encode(text)
-    tokens = tokenizer.convert_ids_to_tokens(ids)
+    # Tokenize
+    ids = mt_tokenizer.encode(input_text)
+    tokens = mt_tokenizer.convert_ids_to_tokens(ids)
 
+    # Translate with CTranslate2
     results = translator.translate_batch(
         [tokens],
         target_prefix=[["eng_Latn"]],
-        beam_size=1
+        beam_size=4
     )
 
     out = results[0].hypotheses[0]
 
+    # Strip target prefix token if present
     if out[0] == "eng_Latn":
         out = out[1:]
 
-    ids = tokenizer.convert_tokens_to_ids(out)
-
-    return tokenizer.decode(ids, skip_special_tokens=True)
-
+    ids = mt_tokenizer.convert_tokens_to_ids(out)
+    return mt_tokenizer.decode(ids, skip_special_tokens=True)
 
 # ====================================
 # PIPELINE
 # ====================================
 def pipeline(audio):
-
     try:
         efik = transcribe(audio)
         eng = translate(efik)
-
         return efik, eng
-
     except Exception as e:
         return f"ERROR: {str(e)}", ""
 
-
 # ====================================
 # UI
 # ====================================
@@ -171,7 +168,7 @@ with gr.Blocks() as demo:
         type="numpy"
     )
 
-    btn = gr.Button("🚀 Start")
+    btn = gr.Button("🚀 Translate")
 
     out1 = gr.Textbox(label="Efik Text")
     out2 = gr.Textbox(label="English")