Update app.py

app.py

@@ -1,13 +1,15 @@
 import os
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"
 
-from transformers import MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig
+from transformers import (
+    MBart50TokenizerFast, MBartForConditionalGeneration, AutoConfig,
+    AutoTokenizer, AutoModelForSeq2SeqLM
+)
 import gradio as gr
 
-# ---- Load model & tokenizer ----
+# ---- Load Translation Model ----
 model_name = "Mudasir692/mbart-eng-ur"
 
-# Fix config issue
 config = AutoConfig.from_pretrained(model_name)
 if getattr(config, "early_stopping", None) is None:
     config.early_stopping = True
@@ -15,6 +17,11 @@ if getattr(config, "early_stopping", None) is None:
 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)
+
 # ---- Language mapping ----
 LANG_CODES = {
     "Urdu": "ur_PK",
@@ -39,45 +46,57 @@ LANG_CODES = {
     "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):
     if not text.strip():
-        return "⚠️ Please enter text to translate."
+        return "⚠️ Please enter text to translate.", ""
 
-    # Source language
+    # Step 1: Grammar correction
+    corrected_text = correct_grammar(text)
+
+    # Step 2: Language detection
     if auto_detect:
-        # Very simple heuristic-based detection
-        if any("\u0600" <= ch <= "\u06FF" for ch in text):
+        if any("\u0600" <= ch <= "\u06FF" for ch in corrected_text):
             src_lang = "ur_PK"
-        elif any("\u0900" <= ch <= "\u097F" for ch in text):
+        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 3: Translation
     tgt_lang_code = LANG_CODES.get(target_lang, "ur_PK")
-
     tokenizer.src_lang = src_lang
     tokenizer.tgt_lang = tgt_lang_code
 
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    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
     )
-    output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
-    return output
+    translated_output = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
+
+    return corrected_text, translated_output
 
 # ---- Examples ----
 examples = [
-    ["How are you?", "Urdu", False],
-    ["Where are you going?", "Arabic", False],
-    ["This is my new project.", "Hindi", False],
+    ["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 help me?", "Spanish", False],
+    ["Can you helps me?", "Spanish", False],
 ]
 
 # ---- Gradio Interface ----
@@ -88,28 +107,28 @@ with gr.Blocks(css="""
 
     gr.Markdown("""
     <div style='text-align:center;'>
-        <h2>🌍 Multi-Language Translator (mBART)</h2>
-        <p>Translate between English and 20+ languages using a fine-tuned mBART model.</p>
+        <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>
     </div>
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            text_input = gr.Textbox(label="Enter Text", placeholder="Type your sentence here...", lines=4)
+            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):
-            result_output = gr.Textbox(label="Translation", lines=4)
-            copy_btn = gr.Button("📋 Copy Translation")
+            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=result_output)
-    #copy_btn.click(None, inputs=result_output, outputs=None, _js="(text) => navigator.clipboard.writeText(text)")
+    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)