Update app.py

app.py

@@ -5,14 +5,14 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, MarianMTModel, Ma
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 MODEL_OPTIONS = [
-    "Helsinki-NLP (Acceptable)", # Round-trip OPUS-MT en→es→en
-    "FLAN-T5-base (Google poor results)"
+    "Helsinki-NLP (Tira ondo)", # Round-trip OPUS-MT en→es→en
+    "FLAN-T5-base (Google gaizki xamar)"
 ]
 
 # Cache
 CACHE = {}
 
-# --- FLAN loader ---
+# --- FLAN loader (Google-style Euskera correction) ---
 def load_flan():
     if "flan" not in CACHE:
         tok = AutoTokenizer.from_pretrained("google/flan-t5-base")
@@ -26,33 +26,33 @@ def load_flan():
 
 def run_flan(sentence: str) -> str:
     model, tok = load_flan()
-    prompt = f"Correct grammar and rewrite in fluent British English: {sentence}"
+    prompt = f"Euskara zuzen gramatikalki eta idatzi modu naturalean: {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"]
+# --- Euskera round-trip loader ---
+def load_euskera():
+    if "eus" not in CACHE:
+        tok1 = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-eu-es")
+        mdl1 = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-eu-es").to(DEVICE)
+        tok2 = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-es-eu")
+        mdl2 = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-es-eu").to(DEVICE)
+        CACHE["eus"] = (mdl1, tok1, mdl2, tok2)
+    return CACHE["eus"]
 
 def run_roundtrip(sentence: str) -> str:
-    mdl1, tok1, mdl2, tok2 = load_marian()
-    # English → Spanish
+    mdl1, tok1, mdl2, tok2 = load_euskera()
+    # Euskera → 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
+    # Spanish → Euskera
     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()
+    eu_tokens = mdl2.generate(**inputs2, max_length=128, num_beams=4)
+    euskera = tok2.decode(eu_tokens[0], skip_special_tokens=True)
+    return euskera.strip()
 
 # --- Dispatcher ---
 def polish(sentence: str, choice: str) -> str:
@@ -66,8 +66,8 @@ def polish(sentence: str, choice: str) -> str:
         return "Unknown option."
 
 # --- Gradio UI ---
-with gr.Blocks(title="HizkuntzLagun: English Fixer (CPU enabled)") as demo:
-    gr.Markdown("### HizkuntzLagun: English Fixer\n")
+with gr.Blocks(title="HizkuntzLagun: Euskera Fixer (CPU enabled)") as demo:
+    gr.Markdown("### HizkuntzLagun: Euskera Fixer\n")
     gr.Markdown(
         """
         > ⚡ **Note:**  
@@ -76,11 +76,11 @@ with gr.Blocks(title="HizkuntzLagun: English Fixer (CPU enabled)") as demo:
         > Expect quick corrections, not deep grammar analysis.
         > Drop in anytime — a quick fix a day keeps awkward grammar away.
         """)
-    inp = gr.Textbox(lines=3, label="Input (English) E.g. She go tomorrow buy two bread.", placeholder="Type an English sentence to correct.")
-    choice = gr.Dropdown(choices=MODEL_OPTIONS, value="Helsinki-NLP", label="Method")
-    btn = gr.Button("Oxford grammar polish")
-    out = gr.Textbox(label="Output")
+    inp = gr.Textbox(lines=3, label="Sarrera (Euskara)", placeholder="Idatzi zuzentzeko esaldi bat...")
+    choice = gr.Dropdown(choices=MODEL_OPTIONS, value="eu_sp_eu", label="Metodoa")
+    btn = gr.Button("Euskara zuzen")
+    out = gr.Textbox(label="Irteera")
     btn.click(polish, inputs=[inp, choice], outputs=out)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()