Update app.py

app.py

@@ -1,6 +1,7 @@
 import requests
 import gradio as gr
 from dotenv import load_dotenv
+from transformers import MBart50TokenizerFast
 import os
 
 # Load environment variables
@@ -8,24 +9,35 @@ load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 headers = {"Authorization": f"Bearer {HF_TOKEN}"}
 
-# Language to ISO 639-3 codes (used for NLLB-200)
+# Correct mBART-50 language codes
 LANGUAGES = {
-    "English → Afrikaans": "afr",
-    "English → Xhosa": "xho",
-    "English → Zulu": "zul",
-    "English → Sesotho": "sot",
-    "English → Tswana": "tsn",
-    "English → Northern Sotho": "nso",
-    "English → Swati": "ssw",
-    "English → Tsonga": "tso",
-    "English → Venda": "ven",
+    "English → Afrikaans": "af_ZA",
+    "English → Xhosa": "xh_ZA",
+    "English → Zulu": "zu_ZA",
+    "English → Sesotho": "st_ZA",  # Southern Sotho
+    "English → Tswana": "tn_ZA",
+    # The following are *not officially* supported by mBART-50 and may raise errors
+    # You can remove them if not working
+    # "English → Northern Sotho": "nso_ZA",
+    # "English → Swati": "ss_ZA",
+    # "English → Tsonga": "ts_ZA",
+    # "English → Venda": "ve_ZA",
 }
 
 MODEL_NAME = "facebook/mbart-large-50-many-to-many-mmt"
 API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
 
+# Load tokenizer to get language token IDs
+tokenizer = MBart50TokenizerFast.from_pretrained(MODEL_NAME)
+
+
+def get_token_id(lang_code):
+    """Return the forced_bos_token_id for the target language."""
+    return tokenizer.lang_code_to_id[lang_code]
+
 
 def query(payload):
+    """Send the translation request to the Hugging Face API."""
     response = requests.post(API_URL, headers=headers, json=payload)
 
     if response.status_code != 200:
@@ -40,10 +52,17 @@ def query(payload):
 
 
 def translate(input_text, language_label):
-    language_code = LANGUAGES[language_label]
-    formatted_input = f">>{language_code}<< {input_text}"
+    """Main translation function."""
+    target_lang_code = LANGUAGES[language_label]
+    token_id = get_token_id(target_lang_code)
+
+    payload = {
+        "inputs": input_text,
+        "parameters": {"forced_bos_token_id": token_id},
+        "options": {"wait_for_model": True},
+    }
 
-    response = query({"inputs": formatted_input, "options": {"wait_for_model": True}})
+    response = query(payload)
 
     if "error" in response:
         return f"Error: {response['error']}"
@@ -51,6 +70,7 @@ def translate(input_text, language_label):
     return response[0]["translation_text"]
 
 
+# Gradio UI
 translator = gr.Interface(
     fn=translate,
     inputs=[
@@ -59,7 +79,7 @@ translator = gr.Interface(
     ],
     outputs=gr.Textbox(label="Translation"),
     title="Translademia",
-    description="Translate English text to South African languages using Meta's NLLB-200 model.",
+    description="Translate English text to South African languages using Meta's mBART-50 model.",
 )
 
 translator.launch()