Update app.py

app.py

@@ -47,44 +47,67 @@
 
 # iface.launch()
 
-from transformers import MarianMTModel, MarianTokenizer, pipeline
+# from transformers import MarianMTModel, MarianTokenizer, pipeline
+# import gradio as gr
+
+# # Define supported models for South African languages
+# language_models = {
+#     "Afrikaans": "Helsinki-NLP/opus-mt-en-af",
+#     "Zulu": "Helsinki-NLP/opus-mt-en-zu",
+#     "Xhosa": "Helsinki-NLP/opus-mt-en-xh",
+#     "Sesotho": "Helsinki-NLP/opus-mt-en-st",
+#     "Setswana": "Helsinki-NLP/opus-mt-en-tn",
+# }
+
+
+# # Translation function
+# def translate(text, target_language):
+#     model_name = language_models[target_language]
+#     tokenizer = MarianTokenizer.from_pretrained(model_name)
+#     model = MarianMTModel.from_pretrained(model_name)
+
+#     # Setup pipeline
+#     translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
+
+#     # Translate
+#     result = translation_pipeline(text)
+#     return result[0]["translation_text"]
+
+
+# # Build Gradio interface
+# interface = gr.Interface(
+#     fn=translate,
+#     inputs=[
+#         gr.Textbox(label="Enter English Text"),
+#         gr.Dropdown(choices=list(language_models.keys()), label="Translate to"),
+#     ],
+#     outputs="text",
+#     title="African Language Translator",
+#     description="Translate English text into Afrikaans, Zulu, Xhosa, Sesotho or Setswana",
+# )
+
+# # Launch the app
+# interface.launch()
+
+
 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 tokenizer and model locally
+tokenizer = MarianTokenizer.from_pretrained(model_name)
+model = MarianMTModel.from_pretrained(model_name)
+
+
+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)
+
 
-# Define supported models for South African languages
-language_models = {
-    "Afrikaans": "Helsinki-NLP/opus-mt-en-af",
-    "Zulu": "Helsinki-NLP/opus-mt-en-zu",
-    "Xhosa": "Helsinki-NLP/opus-mt-en-xh",
-    "Sesotho": "Helsinki-NLP/opus-mt-en-st",
-    "Setswana": "Helsinki-NLP/opus-mt-en-tn",
-}
-
-
-# Translation function
-def translate(text, target_language):
-    model_name = language_models[target_language]
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-
-    # Setup pipeline
-    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
-
-    # Translate
-    result = translation_pipeline(text)
-    return result[0]["translation_text"]
-
-
-# Build Gradio interface
-interface = gr.Interface(
-    fn=translate,
-    inputs=[
-        gr.Textbox(label="Enter English Text"),
-        gr.Dropdown(choices=list(language_models.keys()), label="Translate to"),
-    ],
-    outputs="text",
-    title="African Language Translator",
-    description="Translate English text into Afrikaans, Zulu, Xhosa, Sesotho or Setswana",
-)
-
-# Launch the app
-interface.launch()
+# Gradio interface
+gr.Interface(
+    fn=translate, inputs="text", outputs="text", title="English to Zulu Translator"
+).launch()