Update app.py

app.py

@@ -1,34 +1,81 @@
+# app.py
+# Requirements: transformers, torch, sentencepiece, sacremoses, gradio
+
+import torch
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
-
-# 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")
-
-# 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) -> 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)
-
-    # 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 english
-
-demo = gr.Interface(
-    fn=polish,
-    inputs=gr.Textbox(lines=2, placeholder="Enter a sentence in English..."),
-    outputs=gr.Textbox(label="Corrected English"),
-    title="Round-trip Grammar Polisher",
-    description="Uses Helsinki-NLP MarianMT models (en→es→en) to smooth and correct English sentences."
-)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, MarianMTModel, MarianTokenizer
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+MODEL_OPTIONS = [
+    "FLAN-T5-base (Google en→en)",
+    "Round-trip OPUS-MT en→es→en (Helsinki-NLP)"
+]
+
+# Cache
+CACHE = {}
+
+# --- FLAN loader ---
+def load_flan():
+    if "flan" not in CACHE:
+        tok = AutoTokenizer.from_pretrained("google/flan-t5-base")
+        mdl = AutoModelForSeq2SeqLM.from_pretrained(
+            "google/flan-t5-base",
+            low_cpu_mem_usage=True,
+            torch_dtype="auto"
+        ).to(DEVICE)
+        CACHE["flan"] = (mdl, tok)
+    return CACHE["flan"]
+
+def run_flan(sentence: str) -> str:
+    model, tok = load_flan()
+    prompt = f"Correct grammar and rewrite in fluent British English: {sentence}"
+    inputs = tok(prompt, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        out = model.generate(**inputs, max_new_tokens=96, num_beams=4)
+    return tok.decode(out[0], skip_special_tokens=True).strip()
+
+# --- Marian round-trip loader ---
+def load_marian():
+    if "en_es" not in CACHE:
+        tok1 = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
+        mdl1 = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-es").to(DEVICE)
+        tok2 = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-es-en")
+        mdl2 = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-es-en").to(DEVICE)
+        CACHE["en_es"] = (mdl1, tok1, mdl2, tok2)
+    return CACHE["en_es"]
+
+def run_roundtrip(sentence: str) -> str:
+    mdl1, tok1, mdl2, tok2 = load_marian()
+    # English → Spanish
+    inputs = tok1(sentence, return_tensors="pt").to(DEVICE)
+    es_tokens = mdl1.generate(**inputs, max_length=128, num_beams=4)
+    spanish = tok1.decode(es_tokens[0], skip_special_tokens=True)
+    # Spanish → English
+    inputs2 = tok2(spanish, return_tensors="pt").to(DEVICE)
+    en_tokens = mdl2.generate(**inputs2, max_length=128, num_beams=4)
+    english = tok2.decode(en_tokens[0], skip_special_tokens=True)
+    return english.strip()
+
+# --- Dispatcher ---
+def polish(sentence: str, choice: str) -> str:
+    if not sentence.strip():
+        return ""
+    if choice.startswith("FLAN"):
+        return run_flan(sentence)
+    elif choice.startswith("Round-trip"):
+        return run_roundtrip(sentence)
+    else:
+        return "Unknown option."
+
+# --- Gradio UI ---
+with gr.Blocks(title="English Grammar Polisher") as demo:
+    gr.Markdown("### English Grammar Polisher\nChoose FLAN-T5 (Google) or OPUS-MT round-trip (Helsinki-NLP).")
+    inp = gr.Textbox(lines=3, label="Input (English)", placeholder="Type a sentence…")
+    choice = gr.Dropdown(choices=MODEL_OPTIONS, value="FLAN-T5-base (Google en→en)", label="Method")
+    btn = gr.Button("Polish")
+    out = gr.Textbox(label="Output")
+    btn.click(polish, inputs=[inp, choice], outputs=out)
 
 if __name__ == "__main__":
     demo.launch()