Update app.py

app.py

@@ -47,59 +47,44 @@
 
 # iface.launch()
 
-
-import requests
+from transformers import MarianMTModel, MarianTokenizer, pipeline
 import gradio as gr
-from dotenv import load_dotenv
-import os
-
-# Load env
-load_dotenv()
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-model_name = "Helsinki-NLP/opus-mt-en-nso"
-API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
-headers = {"Authorization": f"Bearer {HF_TOKEN}", "Accept": "application/json"}
-
-
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-
-    if response.status_code == 200:
-        try:
-            return response.json()
-        except Exception as e:
-            return {"error": f"JSON decode error: {str(e)}"}
-    else:
-        # Show error message from HF API or HTTP code
-        try:
-            err = response.json()
-        except:
-            err = response.text
-        return {"error": f"HTTP {response.status_code}: {err}"}
 
+# 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",
+}
 
-def translate(input_text):
-    if not input_text.strip():
-        return "Please enter text to translate."
 
-    response = query({"inputs": input_text, "options": {"wait_for_model": True}})
+# 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)
 
-    if "error" in response:
-        return f"Error: {response['error']}"
+    # Setup pipeline
+    translation_pipeline = pipeline("translation", model=model, tokenizer=tokenizer)
 
-    try:
-        translation = response[0]["translation_text"]
-        return translation
-    except (KeyError, IndexError):
-        return f"Unexpected response format: {response}"
+    # Translate
+    result = translation_pipeline(text)
+    return result[0]["translation_text"]
 
 
-translator = gr.Interface(
+# Build Gradio interface
+interface = gr.Interface(
     fn=translate,
-    inputs=[gr.Textbox(label="Input Text", placeholder="Input Text To Be Translated")],
-    outputs=gr.Textbox(label="Translation"),
-    title="Translademia",
+    inputs=[
+        gr.Textbox(label="Enter English Text"),
+        gr.Dropdown(choices=list(language_models.keys()), label="Translate to"),
+    ],
+    outputs="text",
+    title="🇿🇦 South African Language Translator",
+    description="Translate English text into Afrikaans, Zulu, Xhosa, Sesotho or Setswana",
 )
 
-translator.launch()
+# Launch the app
+interface.launch()