Fix MT error by removing clean_up_tokenization_spaces and address STT warnings

app.py

@@ -93,7 +93,7 @@ def load_models_task():
             logger.info("Loading NLLB-200-distilled-600M model...")
             model_status["mt"] = "loading"
             mt_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
-            mt_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M", clean_up_tokenization_spaces=True)
+            mt_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
             mt_model.to(device)
             logger.info("MT model loaded successfully")
             model_status["mt"] = "loaded"
@@ -111,7 +111,7 @@ def load_models_task():
             logger.info("Loading MMS-TTS model for Tagalog...")
             model_status["tts"] = "loading"
             tts_model = VitsModel.from_pretrained("facebook/mms-tts-tgl")
-            tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-tgl", clean_up_tokenization_spaces=True)
+            tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-tgl")
             tts_model.to(device)
             logger.info("TTS model loaded successfully")
             model_status["tts"] = "loaded"
@@ -121,7 +121,7 @@ def load_models_task():
             try:
                 logger.info("Falling back to MMS-TTS English model...")
                 tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-                tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng", clean_up_tokenization_spaces=True)
+                tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
                 tts_model.to(device)
                 logger.info("Fallback TTS model loaded successfully")
                 model_status["tts"] = "loaded (fallback)"
@@ -189,7 +189,7 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
         logger.info(f"Loading MMS-TTS model for {target_code}...")
         from transformers import VitsModel, AutoTokenizer
         tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{target_code}")
-        tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{target_code}", clean_up_tokenization_spaces=True)
+        tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{target_code}")
         device = "cuda" if torch.cuda.is_available() else "cpu"
         tts_model.to(device)
         logger.info(f"TTS model updated to {target_code}")
@@ -199,7 +199,7 @@ async def update_languages(source_lang: str = Form(...), target_lang: str = Form
         try:
             logger.info("Falling back to MMS-TTS English model...")
             tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-            tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng", clean_up_tokenization_spaces=True)
+            tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
             tts_model.to(device)
             logger.info("Fallback TTS model loaded successfully")
             model_status["tts"] = "loaded (fallback)"
@@ -235,7 +235,7 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
             target_nllb_code = NLLB_LANGUAGE_CODES[target_code]
             mt_tokenizer.src_lang = source_nllb_code
             device = "cuda" if torch.cuda.is_available() else "cpu"
-            inputs = mt_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+            inputs = mt_tokenizer(text, return_tensors="pt").to(device)
             with torch.no_grad():
                 generated_tokens = mt_model.generate(
                     **inputs,
@@ -254,7 +254,7 @@ async def translate_text(text: str = Form(...), source_lang: str = Form(...), ta
     output_audio = None
     if model_status["tts"].startswith("loaded") and tts_model is not None and tts_tokenizer is not None:
         try:
-            inputs = tts_tokenizer(translated_text, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+            inputs = tts_tokenizer(translated_text, return_tensors="pt").to(device)
             with torch.no_grad():
                 output = tts_model(**inputs)
             speech = output.waveform.cpu().numpy().squeeze()
@@ -322,7 +322,11 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         inputs = stt_processor(waveform, sampling_rate=16000, return_tensors="pt").to(device)
         logger.info("Audio processed, generating transcription...")
         with torch.no_grad():
-            generated_ids = stt_model.generate(**inputs)
+            generated_ids = stt_model.generate(
+                **inputs,
+                language="en",  # Explicitly set language to English
+                return_attention_mask=True  # Generate attention mask
+            )
             transcription = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         logger.info(f"Transcription completed: {transcription}")
 
@@ -335,7 +339,7 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
                 source_nllb_code = NLLB_LANGUAGE_CODES[source_code]
                 target_nllb_code = NLLB_LANGUAGE_CODES[target_code]
                 mt_tokenizer.src_lang = source_nllb_code
-                inputs = mt_tokenizer(transcription, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+                inputs = mt_tokenizer(transcription, return_tensors="pt").to(device)
                 with torch.no_grad():
                     generated_tokens = mt_model.generate(
                         **inputs,
@@ -353,7 +357,7 @@ async def translate_audio(audio: UploadFile = File(...), source_lang: str = Form
         # Step 3: Convert translated text to speech (TTS)
         if model_status["tts"].startswith("loaded") and tts_model is not None and tts_tokenizer is not None:
             try:
-                inputs = tts_tokenizer(translated_text, return_tensors="pt", clean_up_tokenization_spaces=True).to(device)
+                inputs = tts_tokenizer(translated_text, return_tensors="pt").to(device)
                 with torch.no_grad():
                     output = tts_model(**inputs)
                 speech = output.waveform.cpu().numpy().squeeze()