summarization.py

"""Built-in summarization task.

This task combines lexical overlap, heuristic grounding checks, and rubric-based
LLM judgment so the final score captures both surface fidelity and deeper summary
quality.
"""

from __future__ import annotations

import re
from pathlib import Path

from eval_harness.providers.base import ModelProvider
from eval_harness.scorers.base import Scorer
from eval_harness.scorers.composite import CompositeScorer
from eval_harness.scorers.llm_judge import COHERENCE, FAITHFULNESS, LLMJudgeScorer
from eval_harness.scorers.rouge import RougeScorer
from eval_harness.templates import PromptTemplate
from eval_harness.types import EvalSample, ScorerOutput

from eval_harness.tasks import TaskDefinition

REPO_ROOT = Path(__file__).resolve().parents[2]
SUMMARY_PATTERN = re.compile(r"<summary>\s*(.*?)\s*</summary>", re.IGNORECASE | re.DOTALL)
STOPWORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "for",
    "from",
    "in",
    "is",
    "it",
    "of",
    "on",
    "or",
    "that",
    "the",
    "to",
    "with",
}
TASK_SYSTEM_PROMPT = "You are a careful summarizer. Preserve the main claim, key evidence, and caveats without adding unsupported facts."


def _content_tokens(text: str) -> set[str]:
    """Extract non-trivial word tokens.

    Args:
        text: Input text.

    Returns:
        Set of lowercased content tokens.

    Example:
        >>> sorted(_content_tokens('The city approved the pilot.'))
        ['approved', 'city', 'pilot']
    """

    return {
        token.lower()
        for token in re.findall(r"\w+", text)
        if token.lower() not in STOPWORDS and len(token) > 2
    }


class FaithfulnessScorer(Scorer):
    """Heuristic scorer that checks whether summary content is grounded in the source."""

    async def score(self, sample: EvalSample, model_output: str) -> ScorerOutput:
        """Score how well summary claims are supported by the source document.

        Args:
            sample: Summarization sample where `sample.input` is the source document.
            model_output: Summary produced by the model.

        Returns:
            Grounding score based on token coverage.

        Example:
            This heuristic is intentionally conservative and complements the LLM judge.
        """

        source_tokens = _content_tokens(sample.input)
        summary_sentences = [sentence.strip() for sentence in re.split(r"(?<=[.!?])\s+", model_output) if sentence.strip()]
        if not summary_sentences:
            return ScorerOutput(
                scorer_name=self.name,
                score=0.0,
                reasoning="Empty summary cannot be faithful to the source.",
                metadata={},
            )
        sentence_scores: list[float] = []
        unsupported_tokens: list[str] = []
        for sentence in summary_sentences:
            tokens = _content_tokens(sentence)
            if not tokens:
                sentence_scores.append(1.0)
                continue
            grounded = tokens & source_tokens
            sentence_scores.append(len(grounded) / len(tokens))
            unsupported_tokens.extend(sorted(tokens - source_tokens))
        score = sum(sentence_scores) / len(sentence_scores)
        reasoning = (
            f"Average sentence grounding score is {score:.3f}. Unsupported tokens include: "
            f"{', '.join(unsupported_tokens[:10]) or 'none'}."
        )
        return ScorerOutput(
            scorer_name=self.name,
            score=score,
            reasoning=reasoning,
            metadata={"sentence_scores": [round(value, 3) for value in sentence_scores]},
        )


class SummarizationTask:
    """Factory for the built-in summarization task."""

    @staticmethod
    def summary_extractor(raw_output: str, sample: EvalSample) -> str | None:
        """Extract the `<summary>` section when present.

        Args:
            raw_output: Raw model output.
            sample: Sample associated with the output.

        Returns:
            Extracted summary or the raw output.

        Example:
            >>> SummarizationTask.summary_extractor('<summary>Hello.</summary>', EvalSample(id='1', input='x', expected_output='y', metadata={}, tags=[]))
            'Hello.'
        """

        del sample
        match = SUMMARY_PATTERN.search(raw_output)
        if match:
            return match.group(1).strip()
        return raw_output.strip() or None

    @classmethod
    def build(
        cls,
        judge_model_id: str | None = None,
        judge_provider: ModelProvider | None = None,
    ) -> TaskDefinition:
        """Build the summarization task definition.

        Args:
            judge_model_id: Optional override for the judge model identifier.
            judge_provider: Optional injected judge provider for testing.

        Returns:
            Fully configured task definition.

        Example:
            >>> SummarizationTask.build().name
            'summarization'
        """

        template = PromptTemplate.from_file(
            REPO_ROOT / "templates" / "summarization.jinja2",
            xml_output_tag="summary",
        )
        scorer = CompositeScorer(
            scorers=[
                RougeScorer(primary_metric="rougeL", weight=0.35),
                FaithfulnessScorer(weight=0.25),
                LLMJudgeScorer(
                    rubric=f"{COHERENCE}\n\n{FAITHFULNESS}",
                    judge_model_id=judge_model_id,
                    judge_provider=judge_provider,
                    weight=0.40,
                ),
            ],
            mode="weighted_average",
        )
        return TaskDefinition(
            name="summarization",
            description="Document summarization with ROUGE-L, grounding, and LLM judge scoring.",
            template=template,
            scorer=scorer,
            system_prompt=TASK_SYSTEM_PROMPT,
            output_parser=cls.summary_extractor,
            few_shot_examples=[],
        )


__all__ = ["FaithfulnessScorer", "SummarizationTask"]