evo_plugin_example.py

"""
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.
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM


class _HFGenerator:
    def __init__(self, model_name: str = "distilgpt2"):
        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

    @torch.no_grad()
    def generate(self, prompt: str, max_new_tokens: int = 200, temperature: float = 0.4) -> str:
        inputs = self.tok(prompt, return_tensors="pt").to(self.device)
        out = self.model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            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()


def load_model():
    """
    (Objective) Entry-point used by evo_inference.py.
    """
    return _HFGenerator()