src/flashcard_generator/summarizer.py

from __future__ import annotations

from functools import lru_cache

from .text_processing import first_sentences, token_count


class Summarizer:
    def __init__(self, model_name: str = "google/flan-t5-small", use_model: bool = True):
        self.model_name = model_name
        self.use_model = use_model

    def summarize(self, text: str) -> str:
        if not text.strip():
            return ""
        if not self.use_model:
            return self._fallback_summary(text)

        try:
            summarizer = _load_pipeline(self.model_name)
            max_length = min(180, max(60, token_count(text) // 2))
            result = summarizer(
                f"summarize: {text}",
                max_length=max_length,
                min_length=min(40, max_length - 10),
                do_sample=False,
            )
            summary = result[0]["summary_text"].strip()
            return summary or self._fallback_summary(text)
        except Exception:
            return self._fallback_summary(text)

    @staticmethod
    def _fallback_summary(text: str) -> str:
        return first_sentences(text, limit=4) or text[:800]


@lru_cache(maxsize=2)
def _load_pipeline(model_name: str):
    from transformers import pipeline

    return pipeline("summarization", model=model_name)