Update app.py

app.py

@@ -1,80 +1,73 @@
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
-from transformers import (
-    MBart50TokenizerFast,
-    MBartForConditionalGeneration,
-    AutoConfig
-)
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig
 import gradio as gr
 
-# ---- Load model and tokenizer ----
+# ---- Load model & tokenizer ----
 model_name = "Mudasir692/mbart-eng-ur"
 
+# Fix config issue
 config = AutoConfig.from_pretrained(model_name)
-
-# Fix missing or invalid parameters
-if getattr(config, "early_stopping", None) is None:
+if config.early_stopping is None:
     config.early_stopping = True
-if getattr(config, "max_length", None) is None:
-    config.max_length = 128  # ✅ set a safe limit
-if getattr(config, "num_beams", None) is None:
-    config.num_beams = 4
 
-tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang="en_XX", tgt_lang="ur_PK")
+tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
 model = MBartForConditionalGeneration.from_pretrained(model_name, config=config)
 
+# ---- Language mapping ----
+LANG_CODES = {
+    "Urdu": "ur_PK",
+    "Arabic": "ar_AR",
+    "Hindi": "hi_IN",
+}
+
 # ---- Translation function ----
-def translate_to_urdu(text):
+def translate_text(text, target_lang):
     if not text.strip():
         return "Please enter some English text."
 
+    tgt_lang_code = LANG_CODES.get(target_lang, "ur_PK")
+    tokenizer.src_lang = "en_XX"
+    tokenizer.tgt_lang = tgt_lang_code
+
     inputs = tokenizer(text, return_tensors="pt", padding=True)
-    translated_tokens = model.generate(
-        **inputs,
-        max_length=128,        # ✅ explicitly set again
-        num_beams=4,
-        early_stopping=True
-    )
-    urdu_output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return urdu_output
+    translated_tokens = model.generate(**inputs)
+    output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    return output
 
-# ---- Gradio Interface ----
+# ---- Examples ----
 examples = [
-    ["How are you?"],
-    ["Today is a beautiful day."],
-    ["Where are you going?"],
-    ["I am learning Artificial Intelligence."],
-    ["Thank you very much!"]
+    ["How are you?", "Urdu"],
+    ["Where are you going?", "Arabic"],
+    ["This is my new project.", "Hindi"],
 ]
 
+# ---- Gradio Interface ----
 app = gr.Interface(
-    fn=translate_to_urdu,
-    inputs=gr.Textbox(label="Enter English Text", placeholder="Type your English sentence here...", lines=2),
-    outputs=gr.Textbox(label="Urdu Translation", lines=2),
+    fn=translate_text,
+    inputs=[
+        gr.Textbox(label="Enter English Text", placeholder="Type your English sentence here...", lines=2),
+        gr.Dropdown(["Urdu", "Arabic", "Hindi"], label="Select Target Language", value="Urdu")
+    ],
+    outputs=gr.Textbox(label="Translation", lines=2),
     examples=examples,
-    title="🕌 English → Urdu Translator",
+    title="🌍 Multi-Language Translator",
     description="""
     <div style='text-align:center;'>
-        <h3>Translate English sentences into Urdu using a fine-tuned mBART model.</h3>
+        <h3>Translate English sentences into Urdu, Arabic, or Hindi using a fine-tuned mBART model.</h3>
         <p style='color:gray;'>Built by <b>Khurram Basharat</b> — powered by Transformers & Gradio.</p>
-        <p><i>⏳ Please wait a few seconds for the model to load when you first open the app.</i></p>
+        <p><i>⏳ The model loads on first use, please wait a few seconds.</i></p>
     </div>
     """,
-    theme="soft",
     css="""
         body {
-            background: linear-gradient(135deg, #d1f2eb, #a3e4d7, #d6eaf8);
+            background: linear-gradient(to bottom right, #f1f2f6, #dff9fb);
             font-family: 'Segoe UI', sans-serif;
-            background-attachment: fixed;
         }
         .gr-button-primary {
             background-color: #1e3799 !important;
             color: white !important;
-            border-radius: 8px !important;
-        }
-        h1, h3 {
-            color: #0a3d62 !important;
         }
     """,
 )