Update app.py

app.py

@@ -1,26 +1,20 @@
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
-from transformers import (
-    MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig,
-    AutoTokenizer, AutoModelForSeq2SeqLM
-)
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig, pipeline
 import gradio as gr
 
-# ---- Load Translation Model ----
+# ---- Load models ----
 model_name = "Mudasir692/mbart-eng-ur"
-
 config = AutoConfig.from_pretrained(model_name)
-if getattr(config, "early_stopping", None) is None:
+if config.early_stopping is None:
     config.early_stopping = True
 
 tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
 model = MBartForConditionalGeneration.from_pretrained(model_name, config=config)
 
-# ---- Load Grammar Correction Model ----
-grammar_model_name = "vennify/t5-base-grammar-correction"
-grammar_tokenizer = AutoTokenizer.from_pretrained(grammar_model_name)
-grammar_model = AutoModelForSeq2SeqLM.from_pretrained(grammar_model_name)
+# Grammar correction pipeline
+grammar_corrector = pipeline("text2text-generation", model="vennify/t5-base-grammar-correction")
 
 # ---- Language mapping ----
 LANG_CODES = {
@@ -28,107 +22,47 @@ LANG_CODES = {
     "Arabic": "ar_AR",
     "Hindi": "hi_IN",
     "French": "fr_XX",
-    "German": "de_DE",
     "Spanish": "es_XX",
-    "Chinese": "zh_CN",
-    "Italian": "it_IT",
-    "Portuguese": "pt_XX",
-    "Russian": "ru_RU",
-    "Japanese": "ja_XX",
-    "Korean": "ko_KR",
-    "Turkish": "tr_TR",
-    "Persian": "fa_IR",
-    "Bengali": "bn_IN",
-    "Punjabi": "pa_IN",
-    "Pashto": "ps_AF",
-    "Malay": "ms_MY",
-    "Indonesian": "id_ID",
-    "Tamil": "ta_IN"
 }
 
-# ---- Grammar Correction Function ----
-def correct_grammar(text):
-    if not text.strip():
-        return text
-    inputs = grammar_tokenizer.encode(f"fix: {text}", return_tensors="pt", max_length=512, truncation=True)
-    outputs = grammar_model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
-    corrected_text = grammar_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return corrected_text
-
 # ---- Translation function ----
-def translate_text(text, target_lang, auto_detect):
+def translate_text(text, target_lang, correct_grammar):
     if not text.strip():
-        return "⚠️ Please enter text to translate.", ""
-
-    # Step 1: Grammar correction
-    corrected_text = correct_grammar(text)
+        return "Please enter some English text."
 
-    # Step 2: Language detection
-    if auto_detect:
-        if any("\u0600" <= ch <= "\u06FF" for ch in corrected_text):
-            src_lang = "ur_PK"
-        elif any("\u0900" <= ch <= "\u097F" for ch in corrected_text):
-            src_lang = "hi_IN"
-        else:
-            src_lang = "en_XX"
-    else:
-        src_lang = "en_XX"
+    # Step 1: Grammar correction (if enabled)
+    corrected_text = text
+    if correct_grammar:
+        result = grammar_corrector(text, max_length=128, num_beams=4)
+        corrected_text = result[0]['generated_text']
 
-    # Step 3: Translation
+    # Step 2: Translation
     tgt_lang_code = LANG_CODES.get(target_lang, "ur_PK")
-    tokenizer.src_lang = src_lang
+    tokenizer.src_lang = "en_XX"
     tokenizer.tgt_lang = tgt_lang_code
 
-    inputs = tokenizer(corrected_text, return_tensors="pt", padding=True, truncation=True)
-    translated_tokens = model.generate(
-        **inputs,
-        max_length=256,
-        num_beams=5,
-        early_stopping=True
-    )
-    translated_output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-
-    return corrected_text, translated_output
+    inputs = tokenizer(corrected_text, return_tensors="pt", padding=True)
+    translated_tokens = model.generate(**inputs)
+    translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
 
-# ---- Examples ----
-examples = [
-    ["I goes to school every day.", "Urdu", False],
-    ["He dont like this movie.", "Hindi", False],
-    ["This is my new project.", "Arabic", False],
-    ["I love learning new languages.", "French", False],
-    ["Can you helps me?", "Spanish", False],
-]
+    return f"✅ Corrected English: {corrected_text}\n\n🌐 Translation ({target_lang}): {translated_text}"
 
 # ---- Gradio Interface ----
-with gr.Blocks(css="""
-    body {background: linear-gradient(to bottom right, #f7f9fb, #e0f7fa);}
-    .gr-button-primary {background-color: #1e3799 !important; color: white !important;}
-""") as app:
-
-    gr.Markdown("""
+app = gr.Interface(
+    fn=translate_text,
+    inputs=[
+        gr.Textbox(label="Enter English Text", placeholder="Type your English sentence here...", lines=2),
+        gr.Dropdown(["Urdu", "Arabic", "Hindi", "French", "Spanish"], label="Select Target Language", value="Urdu"),
+        gr.Checkbox(label="Correct Grammar Before Translation", value=True)
+    ],
+    outputs=gr.Textbox(label="Output (Corrected + Translated)", lines=4),
+    title="🌍 Smart Multi-Language Translator + Grammar Corrector",
+    description="""
     <div style='text-align:center;'>
-        <h2>🌍 Smart Multi-Language Translator</h2>
-        <h4>✨ Includes Grammar Correction Before Translation ✨</h4>
-        <p>Translate between English and 20+ languages using a fine-tuned mBART model with auto grammar correction.</p>
-        <p style='color:gray;'>Built by <b>Khurram Basharat</b> — powered by Hugging Face & Gradio.</p>
+        <h3>Auto-correct English grammar before translating into multiple languages.</h3>
+        <p style='color:gray;'>Powered by Transformers & Gradio — built by <b>Khurram Basharat</b>.</p>
     </div>
-    """)
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            text_input = gr.Textbox(label="Enter Text", placeholder="Type your English sentence...", lines=4)
-            target_lang = gr.Dropdown(sorted(LANG_CODES.keys()), label="Select Target Language", value="Urdu")
-            auto_detect = gr.Checkbox(label="Auto-detect Source Language", value=False)
-            translate_btn = gr.Button("🌐 Translate")
-
-        with gr.Column(scale=1):
-            corrected_output = gr.Textbox(label="Corrected English Sentence", lines=3)
-            translated_output = gr.Textbox(label="Translated Sentence", lines=3)
-
-    gr.Examples(examples, inputs=[text_input, target_lang, auto_detect])
-
-    # ---- Actions ----
-    translate_btn.click(translate_text, inputs=[text_input, target_lang, auto_detect], outputs=[corrected_output, translated_output])
+    """,
+)
 
-# ---- Launch app ----
-app.launch(server_name="0.0.0.0", server_port=7860)
+app.launch()