generator_engine.py

from __future__ import annotations

from typing import Optional

try:
    from transformers import pipeline
except Exception:
    pipeline = None


class GeneratorEngine:
    def __init__(self, model_name: str = "google/flan-t5-small"):
        self.model_name = model_name
        self.pipe = None
        if pipeline is not None:
            try:
                self.pipe = pipeline("text2text-generation", model=model_name)
            except Exception:
                self.pipe = None

    def available(self) -> bool:
        return self.pipe is not None

    def generate(self, prompt: str, max_new_tokens: int = 96) -> Optional[str]:
        if self.pipe is None:
            return None
        try:
            out = self.pipe(prompt, max_new_tokens=max_new_tokens, do_sample=False)
            if out and isinstance(out, list):
                return str(out[0].get("generated_text", "")).strip()
        except Exception:
            return None
        return None