evo-gov-copilot-mu / evo_plugin_example.py
HemanM's picture
Update evo_plugin_example.py
6a117f9 verified
raw
history blame
1.3 kB
# evo_plugin_example.py — small, instruction-following stand-in generator
# The app will use YOUR evo_plugin.py if present; otherwise it falls back to this.
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
class _HFSeq2SeqGenerator:
def __init__(self, model_name: str = "google/flan-t5-small"):
# CPU is fine for demos; no GPU required on HF Spaces basic CPU.
self.device = torch.device("cpu")
self.tok = AutoTokenizer.from_pretrained(model_name)
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:
# Seq2Seq models generate responses that follow instructions better than GPT-2 here.
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,
)
return self.tok.decode(out[0], skip_special_tokens=True).strip()
def load_model():
# The app calls this to obtain a generator instance.
return _HFSeq2SeqGenerator()