summary_tool/service.py

import os

import torch


MODEL_ID = "google/flan-t5-small"


class SummaryService:
    def __init__(self):
        self.pipe = None
        cpu_count = os.cpu_count() or 1
        torch.set_num_threads(max(1, min(4, cpu_count)))

    def summarize(self, text, style, max_words):
        clean_text = " ".join((text or "").split())
        if not clean_text:
            return "", "Paste text first."

        try:
            prompt = self._build_prompt(clean_text, style, max_words)
            summary = self._run_model(prompt)
            if style == "Bullet Points":
                summary = self._normalize_bullets(summary)
            return summary, f"Generated summary with {MODEL_ID}."
        except Exception as exc:
            return "", f"Summarization failed: {type(exc).__name__}: {exc}"

    def _load_pipeline(self):
        if self.pipe is not None:
            return

        from transformers import pipeline

        self.pipe = pipeline(
            "text2text-generation",
            model=MODEL_ID,
            device=-1,
        )

    def _run_model(self, prompt):
        self._load_pipeline()
        result = self.pipe(
            prompt,
            max_new_tokens=220,
            do_sample=False,
        )
        return (result[0].get("generated_text") or "").strip()

    def _build_prompt(self, text, style, max_words):
        if style == "Short":
            instruction = f"Summarize this text in under {max_words} words using plain language."
        elif style == "Detailed":
            instruction = f"Write a detailed summary in under {max_words} words."
        elif style == "Bullet Points":
            instruction = f"Summarize this text as concise bullet points in under {max_words} words."
        else:
            instruction = f"Write a balanced summary in under {max_words} words."
        return f"{instruction}\n\nText:\n{text}"

    def _normalize_bullets(self, text):
        lines = [line.strip(" -") for line in text.splitlines() if line.strip()]
        if not lines:
            return text
        return "\n".join(f"- {line}" for line in lines[:8])