generator.py

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