Update app.py

app.py

@@ -1,94 +1,69 @@
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, MarianMTModel, MarianTokenizer
-import shutil, os
-shutil.rmtree(os.path.expanduser("~/.cache/huggingface"), ignore_errors=True)
-shutil.rmtree(os.path.expanduser("~/.cache/torch"), ignore_errors=True)
 
-DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Supported models (text-only for now)
+MODEL_OPTIONS = {
+    "Phi-3.5 Mini Instruct": "microsoft/Phi-3.5-mini-instruct",
+    "Phi-3.5 MoE Instruct": "microsoft/Phi-3.5-MoE-instruct",
+    "Phi-3 Mini 4K Instruct": "microsoft/Phi-3-mini-4k-instruct",
+    "Phi-3 Mini 128K Instruct": "microsoft/Phi-3-mini-128k-instruct"
+}
 
-MODEL_OPTIONS = [
-    "Helsinki-NLP (Tira ondo)", # Round-trip OPUS-MT en→es→en
-    "FLAN-T5-base (Google gaizki xamar)"
-]
-# Cache
-CACHE = {}
+# Cache for loaded models
+loaded_models = {}
 
-# --- FLAN loader (Google-style Euskera correction) ---
-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"]
+# Load model/tokenizer on demand
+def load_model(model_id):
+    if model_id not in loaded_models:
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            trust_remote_code=True,
+            torch_dtype=torch.float32
+        )
+        model.eval()
+        loaded_models[model_id] = (tokenizer, model)
+    return loaded_models[model_id]
 
-def run_flan(sentence: str) -> str:
-    model, tok = load_flan()
-    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()
+# Chat function
+def chat_with_model(user_input, model_choice):
+    model_id = MODEL_OPTIONS[model_choice]
+    tokenizer, model = load_model(model_id)
+
+    messages = [{"role": "user", "content": user_input}]
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt"
+    ).to("cpu")
 
-# --- 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"]
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=100,
+            do_sample=False,
+            temperature=0.7,
+            top_p=0.9
+        )
 
-def run_roundtrip(sentence: str) -> str:
-    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 → Euskera
-    inputs2 = tok2(spanish, return_tensors="pt").to(DEVICE)
-    eu_tokens = mdl2.generate(**inputs2, max_length=128, num_beams=4)
-    euskera = tok2.decode(eu_tokens[0], skip_special_tokens=True)
-    return euskera.strip()
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    return response.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("Helsinki"):
-        return run_roundtrip(sentence)
-    else:
-        return "Unknown option."
+# Gradio UI
+with gr.Blocks(title="Phi-3 Instruct Explorer") as demo:
+    gr.Markdown("## 🧠 Phi-3 Instruct Explorer\nSwitch between Phi-3 instruct models and test responses on CPU.")
+    with gr.Row():
+        model_choice = gr.Dropdown(label="Choose a model", choices=list(MODEL_OPTIONS.keys()), value="Phi-3.5 Mini Instruct")
+    with gr.Row():
+        user_input = gr.Textbox(label="Your message", placeholder="Ask me anything...")
+    with gr.Row():
+        output = gr.Textbox(label="Model response")
+    with gr.Row():
+        submit = gr.Button("Generate")
 
-# --- Gradio UI ---
-with gr.Blocks(title="HizkuntzLagun: AI Euskera Zuzendu (CPU enabled)") as demo:
-    gr.Image(
-        value="banner.png",
-        show_label=False,
-        elem_id="banner",
-        height=200
-    )
-    gr.Markdown("### HizkuntzLagun: AI Euskera Zuzedu\n")
-    gr.Markdown(
-        """
-        > ⚡ **Oharra:**  
-        > Tresna honek doako, CPU‑lagunko AI ereduak erabiltzen ditu.  
-        > Azkarra eta eskuragarria izateko diseinatuta dago — ez beti perfektua.  
-        > Zuzenketa azkarrak bai, ez analisi gramatikal sakonak.  
-        > Edozein unetan erabil dezakezu — eguneroko zuzenketa txiki batek saihesten du esaldi traketsen lotsa.
-        """)
-    inp = gr.Textbox(lines=3, label="Idatzi Euskeraz esaldi bat, adibidez Gaur Koldo ikusi nuen.", placeholder="Idatzi Euskeraz esaldi bat...")
-    choice = gr.Dropdown(choices=MODEL_OPTIONS, value="Helsinki-NLP (Tira ondo)", label="Metodoa")
-    btn = gr.Button("Euskera zuzendu")
-    out = gr.Textbox(label="Erantzuna")
-    btn.click(polish, inputs=[inp, choice], outputs=out)
+    submit.click(fn=chat_with_model, inputs=[user_input, model_choice], outputs=output)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()