Update app.py

app.py

@@ -89,25 +89,56 @@
 # # Launch the app
 # interface.launch()
 
-
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
 
-# Pick your language pair
-model_name = "Helsinki-NLP/opus-mt-en-zu"  # English to Zulu
+# Load the tokenizer and model
+model_name = "facebook/nllb-200-distilled-600M"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+
+# Language code map
+lang_map = {
+    "English": "eng_Latn",
+    "Afrikaans": "afr_Latn",
+    "Zulu": "zul_Latn",
+    "Xhosa": "xho_Latn",
+    "French": "fra_Latn",
+    "Spanish": "spa_Latn",
+    "Swahili": "swh_Latn",
+}
+
 
-# Load tokenizer and model locally
-tokenizer = MarianTokenizer.from_pretrained(model_name)
-model = MarianMTModel.from_pretrained(model_name)
+# Translation function
+def translate(text, src_lang, tgt_lang):
+    src_code = lang_map[src_lang]
+    tgt_code = lang_map[tgt_lang]
 
+    tokenizer.src_lang = src_code
+    inputs = tokenizer(text, return_tensors="pt", padding=True)
 
-def translate(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    translated = model.generate(**inputs)
-    return tokenizer.decode(translated[0], skip_special_tokens=True)
+    generated_tokens = model.generate(
+        **inputs, forced_bos_token_id=tokenizer.lang_code_to_id[tgt_code]
+    )
+    translated = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    return translated
 
 
 # Gradio interface
-gr.Interface(
-    fn=translate, inputs="text", outputs="text", title="English to Zulu Translator"
-).launch()
+iface = gr.Interface(
+    fn=translate,
+    inputs=[
+        gr.Textbox(label="Enter text"),
+        gr.Dropdown(
+            choices=list(lang_map.keys()), label="From Language", value="English"
+        ),
+        gr.Dropdown(
+            choices=list(lang_map.keys()), label="To Language", value="Afrikaans"
+        ),
+    ],
+    outputs="text",
+    title="NLLB-200 Custom Language Translator",
+    description="Translate text using Facebook's distilled NLLB-200 model with selectable languages.",
+)
+
+iface.launch()