Update app.py

app.py

@@ -1,59 +1,33 @@
-# app.py
-
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-
-# Translation models for English→English correction
-MODEL_OPTIONS = {
-    "Helsinki-NLP/opus-mt-en-en (light, CPU-friendly)": "Helsinki-NLP/opus-mt-en-en",
-    "facebook/mbart-large-50-many-to-many-mmt (heavier)": "facebook/mbart-large-50-many-to-many-mmt"
-}
+from transformers import MarianMTModel, MarianTokenizer
 
-# Cache loaded pipelines
-loaded_pipelines = {}
+# Load English→Spanish
+en_es_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-es")
+en_es_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
 
-def get_pipeline(model_id: str):
-    if model_id not in loaded_pipelines:
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        model = AutoModelForSeq2SeqLM.from_pretrained(
-            model_id,
-            low_cpu_mem_usage=True,
-            torch_dtype="auto"
-        )
-        pipe = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)
-        # Warm-up
-        _ = pipe("This is a test.", max_length=32)
-        loaded_pipelines[model_id] = pipe
-    return loaded_pipelines[model_id]
+# Load Spanish→English
+es_en_model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-es-en")
+es_en_tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-es-en")
 
-def polish(sentence: str, model_choice: str) -> str:
-    model_id = MODEL_OPTIONS[model_choice]
-    translator = get_pipeline(model_id)
+def polish(sentence: str) -> str:
+    # Step 1: English → Spanish
+    inputs = en_es_tokenizer(sentence, return_tensors="pt", padding=True)
+    translated = en_es_model.generate(**inputs, max_length=128)
+    spanish = en_es_tokenizer.decode(translated[0], skip_special_tokens=True)
 
-    # For mbart we need to set language codes
-    if "mbart" in model_id:
-        inputs = translator.tokenizer(sentence, return_tensors="pt")
-        inputs["forced_bos_token_id"] = translator.tokenizer.lang_code_to_id["en_XX"]
-        out = translator.model.generate(**inputs, max_length=128, num_beams=4)
-        text = translator.tokenizer.decode(out[0], skip_special_tokens=True)
-    else:
-        out = translator(sentence, max_length=128)
-        text = out[0]["translation_text"]
+    # Step 2: Spanish → English
+    inputs2 = es_en_tokenizer(spanish, return_tensors="pt", padding=True)
+    back_translated = es_en_model.generate(**inputs2, max_length=128)
+    english = es_en_tokenizer.decode(back_translated[0], skip_special_tokens=True)
 
-    return text.strip()
+    return english
 
-# Gradio interface
 demo = gr.Interface(
     fn=polish,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter a sentence to correct..."),
-        gr.Dropdown(choices=list(MODEL_OPTIONS.keys()),
-                    value="Helsinki-NLP/opus-mt-en-en (light, CPU-friendly)",
-                    label="Choose Model")
-    ],
+    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence in English..."),
     outputs=gr.Textbox(label="Corrected English"),
-    title="English→English Grammar Polisher",
-    description="Uses translation models (Helsinki-NLP opus-mt-en-en and facebook mbart-large-50) to rewrite English sentences into fluent, corrected English."
+    title="Round-trip Grammar Polisher",
+    description="Uses Helsinki-NLP MarianMT models (en→es→en) to smooth and correct English sentences."
 )
 
 if __name__ == "__main__":