Spaces:
Running on Zero
Running on Zero
File size: 1,848 Bytes
d10de1b 93da3ee d10de1b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | """
Generate text from a prompt using a Hugging Face causal/seq2seq model.
The default model (flan-t5-base) is CPU-friendly for Spaces.
Switch GENERATOR_MODEL in config.py to use Mistral, Gemma, BioMistral, etc.
"""
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
from config import GENERATOR_MODEL, MAX_NEW_TOKENS, TEMPERATURE
_generator = None # lazy singleton
def get_generator():
global _generator
if _generator is None:
print(f"[generator] Loading model: {GENERATOR_MODEL}")
# Detect model type: seq2seq (T5, BART) vs causal (GPT, Mistral, Gemma, Llama)
tokenizer = AutoTokenizer.from_pretrained(GENERATOR_MODEL)
try:
# Try seq2seq first (flan-t5, bart, etc.)
AutoModelForSeq2SeqLM.from_pretrained(GENERATOR_MODEL)
task = "text2text-generation"
except Exception:
task = "text-generation"
_generator = pipeline(
task,
model=GENERATOR_MODEL,
tokenizer=tokenizer,
max_new_tokens=MAX_NEW_TOKENS,
temperature=TEMPERATURE,
do_sample=TEMPERATURE > 0,
)
print(f"[generator] Model ready (task={task}).")
return _generator
def generate_answer(prompt: str) -> str:
"""Run the prompt through the model and return the generated text."""
gen = get_generator()
outputs = gen(prompt)
# Both tasks return a list with one item
result = outputs[0]
if "generated_text" in result:
text = result["generated_text"]
# For causal models the prompt is included in output — strip it
if text.startswith(prompt):
text = text[len(prompt):].strip()
return text
if "summary_text" in result:
return result["summary_text"]
# Fallback key
return str(result)
|