Modified the _initialize_tts_model method to include the clean_up_tokenization_spaces parameter; Added logging configuration to configure the logging level for transformers in app.py

app.py

@@ -24,10 +24,13 @@ from typing import Optional, Tuple, Dict
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from fastapi.responses import JSONResponse
 import tempfile
+import logging
+
+# Configure transformers logging to reduce verbosity
+logging.getLogger("transformers").setLevel(logging.ERROR)
 
 app = FastAPI(title="Talklas API")
 
-# Rest of your code remains the same
 class TalklasTranslator:
     LANGUAGE_MAPPING = {
         "English": "eng",
@@ -79,19 +82,26 @@ class TalklasTranslator:
             self.mt_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
             self.mt_tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
             self.mt_model.to(self.device)
+            print("Loaded NLLB translation model successfully")
         except Exception as e:
             raise RuntimeError(f"MT model initialization failed: {e}")
 
     def _initialize_tts_model(self):
         try:
             self.tts_model = VitsModel.from_pretrained(f"facebook/mms-tts-{self.target_lang}")
-            self.tts_tokenizer = AutoTokenizer.from_pretrained(f"facebook/mms-tts-{self.target_lang}")
+            self.tts_tokenizer = AutoTokenizer.from_pretrained(
+                f"facebook/mms-tts-{self.target_lang}",
+                clean_up_tokenization_spaces=True
+            )
             self.tts_model.to(self.device)
             print(f"Loaded TTS model facebook/mms-tts-{self.target_lang} successfully")
         except Exception:
             print(f"Failed to load facebook/mms-tts-{self.target_lang}, falling back to English TTS")
             self.tts_model = VitsModel.from_pretrained("facebook/mms-tts-eng")
-            self.tts_tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-eng")
+            self.tts_tokenizer = AutoTokenizer.from_pretrained(
+                "facebook/mms-tts-eng",
+                clean_up_tokenization_spaces=True
+            )
             self.tts_model.to(self.device)
             print("Loaded fallback TTS model facebook/mms-tts-eng successfully")
 
@@ -113,18 +123,18 @@ class TalklasTranslator:
             transcription = self.stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return transcription
 
-def translate_text(self, text: str) -> str:
-    source_code = self.NLLB_LANGUAGE_CODES[self.source_lang]
-    target_code = self.NLLB_LANGUAGE_CODES[self.target_lang]
-    self.mt_tokenizer.src_lang = source_code
-    inputs = self.mt_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)
-    with torch.no_grad():
-        generated_tokens = self.mt_model.generate(
-            **inputs,
-            forced_bos_token_id=self.mt_tokenizer.convert_tokens_to_ids(target_code),
-            max_length=448
-        )
-    return self.mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    def translate_text(self, text: str) -> str:
+        source_code = self.NLLB_LANGUAGE_CODES[self.source_lang]
+        target_code = self.NLLB_LANGUAGE_CODES[self.target_lang]
+        self.mt_tokenizer.src_lang = source_code
+        inputs = self.mt_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)
+        with torch.no_grad():
+            generated_tokens = self.mt_model.generate(
+                **inputs,
+                forced_bos_token_id=self.mt_tokenizer.convert_tokens_to_ids(target_code),
+                max_length=448
+            )
+        return self.mt_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
     def text_to_speech(self, text: str) -> Tuple[int, np.ndarray]:
         inputs = self.tts_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)