Update src/streamlit_app.py

src/streamlit_app.py

@@ -43,25 +43,41 @@ def instantiate_translation_model(model_name: str, text: str, src_lg: str, tgt_l
         tgt = "por_Latn" if "pt" in tgt_lg else "kea_Latn"
         pipe = pipeline("translation", model=model_path, token=token, src_lang=src, tgt_lang=tgt)
     elif "m2m100" in model_name:
-        src = "kea_Latn" if "en" in src_lg else "pt"
-        tgt = "por_Latn" if "pt" in tgt_lg else "en"
-        pipe = pipeline("translation", model=model_path, token=token, src_lang=src_lg, tgt_lang=tgt_lg)
+        pipe = pipeline(
+            "translation",
+            model=model_path,
+            tokenizer=model_path,
+            token=token
+        )
+    
+        pipe.tokenizer.src_lang = src_lg
     else:
         # Standard logic for MBart
         pipe = pipeline("translation", model=model_path, token=token, src_lang=src_lg, tgt_lang=tgt_lg)
         
-    result = pipe(text)
+    result = pipe(text, forced_bos_token_id=pipe.tokenizer.get_lang_id(tgt_lg))
     return result[0]["translation_text"]
 
-# --- UI Build Functions ---
 def build_translation_page(model_name):
     st.title(f"🌍 {model_name}: Tradução")
-    
-    # Dynamic language mapping based on model
+
     if "nllb" in model_name:
-        lang_map = {"Português": "por_Latn", "Kabuverdianu": "kea_Latn"}
-    else:
-        lang_map = {"Português": "pt_XX", "Kabuverdianu": "en_XX"} # MBart style
+        lang_map = {
+            "Português": "por_Latn",
+            "Kabuverdianu": "kea_Latn"
+        }
+
+    elif "m2m100" in model_name:
+        lang_map = {
+            "Português": "pt",
+            "Kabuverdianu": "en"  # m2m100 does not support kea
+        }
+
+    else:  # mBART
+        lang_map = {
+            "Português": "pt_XX",
+            "Kabuverdianu": "en_XX"
+        }
 
     col1, col2 = st.columns(2)
     with col1: