Modified the tokenization step to include clean_up_tokenization_spaces=True; Added clean_up_tokenization_spaces=True in the text_to_speech method; Added a print statement to confirm the TTS model is loaded

app.py

@@ -87,10 +87,13 @@ class TalklasTranslator:
             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_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_model.to(self.device)
+            print("Loaded fallback TTS model facebook/mms-tts-eng successfully")
 
     def update_languages(self, source_lang: str, target_lang: str):
         self.source_lang = source_lang
@@ -110,21 +113,21 @@ 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").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").to(self.device)
+        inputs = self.tts_tokenizer(text, return_tensors="pt", clean_up_tokenization_spaces=True).to(self.device)
         with torch.no_grad():
             output = self.tts_model(**inputs)
         speech = output.waveform.cpu().numpy().squeeze()