Update app.py

app.py

@@ -1,39 +1,75 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import os
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
+
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig
 import gradio as gr
 
-# Load model and tokenizer
-model_name = "Helsinki-NLP/opus-mt-en-ur"  # English ↔ Urdu model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# ---- Load model & tokenizer ----
+model_name = "Mudasir692/mbart-eng-ur"
+
+# Fix config issue
+config = AutoConfig.from_pretrained(model_name)
+if config.early_stopping is None:
+    config.early_stopping = True
+
+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",
+}
 
-def translate_text(text):
+# ---- Translation function ----
+def translate_text(text, target_lang):
     if not text.strip():
-        return "Please enter some text to translate."
-    
-    # Tokenize input
-    inputs = tokenizer(text, return_tensors="pt", truncation=True)
-
-    # Generate translation (explicit max_length fix)
-    translated_tokens = model.generate(
-        **inputs,
-        max_length=256,       # Fix for max_length=None error
-        num_beams=5,          # Enables beam search (more accurate)
-        early_stopping=True
-    )
-
-    # Decode output
-    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return translated_text
-
-# Build Gradio Interface
+        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)
+    output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+    return output
+
+# ---- Examples ----
+examples = [
+    ["How are you?", "Urdu"],
+    ["Where are you going?", "Arabic"],
+    ["This is my new project.", "Hindi"],
+]
+
+# ---- Gradio Interface ----
 app = gr.Interface(
     fn=translate_text,
-    inputs=gr.Textbox(lines=3, placeholder="Enter English or Urdu text..."),
-    outputs="text",
-    title="English ↔ Urdu Translator",
-    description="Translate between English and Urdu using a Hugging Face translation model.",
+    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="🌍 Multi-Language Translator",
+    description="""
+    <div style='text-align:center;'>
+        <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>⏳ The model loads on first use, please wait a few seconds.</i></p>
+    </div>
+    """,
+    css="""
+        body {
+            background: linear-gradient(to bottom right, #f1f2f6, #dff9fb);
+            font-family: 'Segoe UI', sans-serif;
+        }
+        .gr-button-primary {
+            background-color: #1e3799 !important;
+            color: white !important;
+        }
+    """,
 )
 
-# Launch app
-if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860)
+app.launch()