src/deberta_ime/benchmark.py

from __future__ import annotations

import math
import random
import re
import time
from collections import Counter, defaultdict
from collections.abc import Callable, Sequence
from dataclasses import dataclass

from .domain import (
    Candidate,
    CandidateScorer,
    CandidateScoringError,
    RerankConfig,
    RerankRequest,
)
from .reranker import Reranker

JAPANESE_RE = re.compile(r"[ぁ-んァ-ヶ一-龯々〆ヵヶ]")


@dataclass(frozen=True, slots=True)
class CorpusToken:
    surface: str
    reading: str
    upos: str


@dataclass(frozen=True, slots=True)
class BenchmarkExample:
    request: RerankRequest
    expected: str


@dataclass(frozen=True, slots=True)
class SplitCoverage:
    eligible_tokens: int
    ambiguous_known_reading: int
    oracle_in_pool: int
    oracle_miss: int


@dataclass(frozen=True, slots=True)
class PreparedSplit:
    examples: tuple[BenchmarkExample, ...]
    coverage: SplitCoverage


@dataclass(frozen=True, slots=True)
class ScoredExample:
    request: RerankRequest
    expected: str
    model_scores: tuple[float, ...] | None
    latency_ms: float
    error_code: str | None = None


@dataclass(frozen=True, slots=True)
class ScoredSplit:
    examples: tuple[ScoredExample, ...]
    total_available: int
    scoring_errors: int
    elapsed_seconds: float

    @property
    def examples_per_second(self) -> float:
        if self.elapsed_seconds == 0.0:
            return 0.0
        return len(self.examples) / self.elapsed_seconds


@dataclass(frozen=True, slots=True)
class SettingMetrics:
    total: int
    baseline_correct: int
    reranked_correct: int
    improved: int
    regressed: int
    both_correct: int
    both_wrong: int
    changed: int

    @property
    def baseline_accuracy(self) -> float:
        return self.baseline_correct / self.total if self.total else 0.0

    @property
    def reranked_accuracy(self) -> float:
        return self.reranked_correct / self.total if self.total else 0.0

    @property
    def absolute_gain(self) -> float:
        return self.reranked_accuracy - self.baseline_accuracy


@dataclass(frozen=True, slots=True)
class SelectedSetting:
    prior_weight: float
    min_margin: float
    metrics: SettingMetrics


@dataclass(frozen=True, slots=True)
class ComparisonOutcome:
    example: ScoredExample
    baseline_prediction: str
    reranked_prediction: str
    baseline_correct: bool
    reranked_correct: bool
    changed: bool
    reason: str


class _PrecomputedScorer:
    def __init__(self, scores: Sequence[float]) -> None:
        self._scores = scores

    def score_candidates(self, request: RerankRequest) -> Sequence[float]:
        return self._scores


def parse_conllu(text: str) -> tuple[tuple[CorpusToken, ...], ...]:
    sentences: list[tuple[CorpusToken, ...]] = []
    current: list[CorpusToken] = []
    for line in text.splitlines():
        if not line:
            if current:
                sentences.append(tuple(current))
                current = []
            continue
        if line.startswith("#"):
            continue
        fields = line.split("\t")
        if len(fields) != 10 or not fields[0].isdigit():
            continue
        reading = ""
        for item in fields[9].split("|"):
            if item.startswith("UnidicInfo="):
                parts = item.removeprefix("UnidicInfo=").split(",")
                lemma_reading = parts[0] if parts else ""
                reading = parts[4] if len(parts) > 4 and parts[4] else lemma_reading
                break
        current.append(CorpusToken(surface=fields[1], reading=reading, upos=fields[3]))
    if current:
        sentences.append(tuple(current))
    return tuple(sentences)


def prepare_examples(
    train_sentences: tuple[tuple[CorpusToken, ...], ...],
    evaluation_sentences: tuple[tuple[CorpusToken, ...], ...],
    *,
    pool_size: int,
) -> PreparedSplit:
    lexicon: dict[tuple[str, str], Counter[str]] = defaultdict(Counter)
    for sentence in train_sentences:
        for token in sentence:
            if token.reading and JAPANESE_RE.search(token.surface):
                lexicon[(token.reading, token.upos)][token.surface] += 1

    eligible_tokens = 0
    ambiguous_known_reading = 0
    oracle_in_pool = 0
    oracle_miss = 0
    examples: list[BenchmarkExample] = []
    for sentence in evaluation_sentences:
        surfaces = tuple(token.surface for token in sentence)
        for target_index, token in enumerate(sentence):
            if not token.reading or not JAPANESE_RE.search(token.surface):
                continue
            eligible_tokens += 1
            counts = lexicon.get((token.reading, token.upos))
            if not counts or len(counts) < 2:
                continue
            ambiguous_known_reading += 1
            pool = sorted(counts.items(), key=lambda item: (-item[1], item[0]))[:pool_size]
            if token.surface not in {surface for surface, _ in pool}:
                oracle_miss += 1
                continue
            oracle_in_pool += 1
            examples.append(
                BenchmarkExample(
                    request=RerankRequest(
                        reading=token.reading,
                        left_context=surfaces[:target_index],
                        right_context=surfaces[target_index + 1 :],
                        candidates=tuple(
                            Candidate(surface=surface, prior_score=math.log1p(count))
                            for surface, count in pool
                        ),
                    ),
                    expected=token.surface,
                )
            )
    return PreparedSplit(
        examples=tuple(examples),
        coverage=SplitCoverage(
            eligible_tokens=eligible_tokens,
            ambiguous_known_reading=ambiguous_known_reading,
            oracle_in_pool=oracle_in_pool,
            oracle_miss=oracle_miss,
        ),
    )


def evaluate_setting(
    examples: Sequence[ScoredExample],
    *,
    prior_weight: float,
    min_margin: float,
) -> SettingMetrics:
    outcomes = compare_examples(
        examples,
        prior_weight=prior_weight,
        min_margin=min_margin,
    )
    return SettingMetrics(
        total=len(examples),
        baseline_correct=sum(outcome.baseline_correct for outcome in outcomes),
        reranked_correct=sum(outcome.reranked_correct for outcome in outcomes),
        improved=sum(
            not outcome.baseline_correct and outcome.reranked_correct for outcome in outcomes
        ),
        regressed=sum(
            outcome.baseline_correct and not outcome.reranked_correct for outcome in outcomes
        ),
        both_correct=sum(
            outcome.baseline_correct and outcome.reranked_correct for outcome in outcomes
        ),
        both_wrong=sum(
            not outcome.baseline_correct and not outcome.reranked_correct
            for outcome in outcomes
        ),
        changed=sum(outcome.changed for outcome in outcomes),
    )


def compare_examples(
    examples: Sequence[ScoredExample],
    *,
    prior_weight: float,
    min_margin: float,
) -> tuple[ComparisonOutcome, ...]:
    outcomes: list[ComparisonOutcome] = []
    for example in examples:
        baseline_prediction = example.request.candidates[0].surface
        if example.model_scores is None:
            reranked_prediction = baseline_prediction
            changed = False
            reason = "scoring_error"
        else:
            result = Reranker(
                _PrecomputedScorer(example.model_scores),
                RerankConfig(prior_weight=prior_weight, min_margin=min_margin),
            ).rerank(example.request)
            reranked_prediction = result.ranked[0].surface
            changed = result.changed
            reason = result.reason
        outcomes.append(
            ComparisonOutcome(
                example=example,
                baseline_prediction=baseline_prediction,
                reranked_prediction=reranked_prediction,
                baseline_correct=baseline_prediction == example.expected,
                reranked_correct=reranked_prediction == example.expected,
                changed=changed,
                reason=reason,
            )
        )
    return tuple(outcomes)


def select_setting(
    examples: Sequence[ScoredExample],
    *,
    prior_weights: Sequence[float],
    min_margins: Sequence[float],
) -> SelectedSetting:
    selected: SelectedSetting | None = None
    selected_key: tuple[int, int, int, float, float] | None = None
    for prior_weight in prior_weights:
        for min_margin in min_margins:
            metrics = evaluate_setting(
                examples,
                prior_weight=prior_weight,
                min_margin=min_margin,
            )
            key = (
                metrics.reranked_correct,
                -metrics.regressed,
                -metrics.changed,
                min_margin,
                prior_weight,
            )
            if selected_key is None or key > selected_key:
                selected_key = key
                selected = SelectedSetting(
                    prior_weight=prior_weight,
                    min_margin=min_margin,
                    metrics=metrics,
                )
    if selected is None:
        raise ValueError("at least one prior weight and margin are required")
    return selected


def mcnemar_exact_p(*, improved: int, regressed: int) -> float:
    discordant = improved + regressed
    if discordant == 0:
        return 1.0
    tail = min(improved, regressed)
    log_probabilities = [
        math.lgamma(discordant + 1)
        - math.lgamma(value + 1)
        - math.lgamma(discordant - value + 1)
        - discordant * math.log(2.0)
        for value in range(tail + 1)
    ]
    largest = max(log_probabilities)
    one_sided = math.exp(largest) * sum(
        math.exp(value - largest) for value in log_probabilities
    )
    return min(1.0, 2.0 * one_sided)


def paired_bootstrap_gain_interval(
    differences: Sequence[int],
    *,
    samples: int,
    seed: int,
) -> tuple[float, float]:
    if not differences:
        return (0.0, 0.0)
    if samples < 1:
        raise ValueError("samples must be positive")
    rng = random.Random(seed)
    count = len(differences)
    gains = sorted(
        sum(differences[rng.randrange(count)] for _ in range(count)) / count
        for _ in range(samples)
    )
    lower_index = int(0.025 * (samples - 1))
    upper_index = math.ceil(0.975 * (samples - 1))
    return gains[lower_index], gains[upper_index]


def score_prepared_split(
    prepared: PreparedSplit,
    scorer: CandidateScorer,
    *,
    limit: int | None = None,
    seed: int = 20260810,
    progress: Callable[[int, int], None] | None = None,
) -> ScoredSplit:
    available = prepared.examples
    if limit is not None and limit < len(available):
        rng = random.Random(seed)
        indexes = sorted(rng.sample(range(len(available)), limit))
        selected = tuple(available[index] for index in indexes)
    else:
        selected = available

    scored_examples: list[ScoredExample] = []
    scoring_errors = 0
    started = time.perf_counter()
    for completed, example in enumerate(selected, start=1):
        row_started = time.perf_counter()
        error_code: str | None = None
        try:
            values = tuple(float(value) for value in scorer.score_candidates(example.request))
            if len(values) != len(example.request.candidates) or not all(
                math.isfinite(value) for value in values
            ):
                raise ValueError("scorer returned invalid scores")
            model_scores: tuple[float, ...] | None = values
        except CandidateScoringError as error:
            model_scores = None
            scoring_errors += 1
            error_code = error.code
        except Exception:
            model_scores = None
            scoring_errors += 1
            error_code = "scorer_exception"
        scored_examples.append(
            ScoredExample(
                request=example.request,
                expected=example.expected,
                model_scores=model_scores,
                latency_ms=(time.perf_counter() - row_started) * 1000.0,
                error_code=error_code,
            )
        )
        if progress is not None:
            progress(completed, len(selected))
    return ScoredSplit(
        examples=tuple(scored_examples),
        total_available=len(available),
        scoring_errors=scoring_errors,
        elapsed_seconds=time.perf_counter() - started,
    )