Update evo_plugin_example.py

evo_plugin_example.py

@@ -1,29 +1,12 @@
-"""
-evo_plugin_example.py
-Step 8: Example text generator plugin (for immediate testing).
-(Objective)
-- Provides the same interface your real Evo plugin will expose.
-- Uses a tiny HF model (distilgpt2) so you can test generative mode now.
-- Replace this later with `evo_plugin.py` wrapping your EvoDecoder/Evo QA.
-
-Real Evo instructions:
-- Create a file `evo_plugin.py` with a `load_model()` that returns an object
-  exposing `generate(prompt: str, max_new_tokens: int, temperature: float) -> str`.
-- The app will auto-prefer `evo_plugin.py` if it exists.
-"""
-
+# evo_plugin_example.py — FLAN-T5 stand-in (better instruction following)
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-
-class _HFGenerator:
-    def __init__(self, model_name: str = "distilgpt2"):
+class _HFSeq2SeqGenerator:
+    def __init__(self, model_name: str = "google/flan-t5-small"):
         self.device = torch.device("cpu")
         self.tok = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
-        # GPT-2 models have no pad token; set to eos to avoid warnings
-        if self.tok.pad_token_id is None and self.tok.eos_token_id is not None:
-            self.tok.pad_token_id = self.tok.eos_token_id
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(self.device).eval()
 
     @torch.no_grad()
     def generate(self, prompt: str, max_new_tokens: int = 200, temperature: float = 0.4) -> str:
@@ -34,15 +17,8 @@ class _HFGenerator:
             do_sample=temperature > 0.0,
             temperature=float(max(0.01, temperature)),
             top_p=0.95,
-            pad_token_id=self.tok.pad_token_id,
         )
-        text = self.tok.decode(out[0], skip_special_tokens=True)
-        # Return only the completion after the prompt to reduce echo
-        return text[len(prompt):].strip()
-
+        return self.tok.decode(out[0], skip_special_tokens=True).strip()
 
 def load_model():
-    """
-    (Objective) Entry-point used by evo_inference.py.
-    """
-    return _HFGenerator()
+    return _HFSeq2SeqGenerator()