src/deberta_ime/benchmark_runner.py

from __future__ import annotations

import json
import platform
import random
import sys
from collections.abc import Callable
from dataclasses import asdict, dataclass, replace
from datetime import UTC, datetime
from importlib.metadata import version
from pathlib import Path
from typing import Any

from .benchmark import (
    ComparisonOutcome,
    PreparedSplit,
    ScoredSplit,
    SettingMetrics,
    compare_examples,
    evaluate_setting,
    mcnemar_exact_p,
    paired_bootstrap_gain_interval,
    parse_conllu,
    prepare_examples,
    score_prepared_split,
    select_setting,
)
from .deberta import MODEL_ID, MODEL_REVISION
from .domain import CandidateScorer
from .ud_gsd import UD_LICENSE, UD_REPOSITORY, UD_REVISION, CorpusArtifact, load_pinned_split


@dataclass(frozen=True, slots=True)
class BenchmarkConfig:
    pool_size: int = 8
    limit_per_split: int | None = None
    seed: int = 20260810
    prior_weights: tuple[float, ...] = (0.0, 0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 4.0)
    min_margins: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0)
    bootstrap_samples: int = 5000
    example_limit: int = 20
    prototype_test_exclusion_count: int = 120
    prototype_test_exclusion_seed: int = 20260811


@dataclass(frozen=True, slots=True)
class BenchmarkRun:
    report: dict[str, Any]
    markdown: str


ProgressCallback = Callable[[str, int, int], None]


def run_benchmark(
    scorer: CandidateScorer,
    *,
    data_dir: Path,
    config: BenchmarkConfig | None = None,
    model_load_seconds: float | None = None,
    progress: ProgressCallback | None = None,
    artifact_loader: Callable[[str, Path], CorpusArtifact] = load_pinned_split,
) -> BenchmarkRun:
    active_config = config or BenchmarkConfig()
    artifacts = {split: artifact_loader(split, data_dir) for split in ("train", "dev", "test")}
    if progress is not None:
        progress("parse", 0, 3)
    train_sentences = parse_conllu(artifacts["train"].text)
    if progress is not None:
        progress("parse", 1, 3)
    dev_sentences = parse_conllu(artifacts["dev"].text)
    if progress is not None:
        progress("parse", 2, 3)
    test_sentences = parse_conllu(artifacts["test"].text)
    if progress is not None:
        progress("parse", 3, 3)

    dev_prepared = prepare_examples(
        train_sentences,
        dev_sentences,
        pool_size=active_config.pool_size,
    )
    test_prepared_all = prepare_examples(
        train_sentences,
        test_sentences,
        pool_size=active_config.pool_size,
    )
    test_prepared, excluded_test_rows = _exclude_prototype_test_rows(
        test_prepared_all,
        count=active_config.prototype_test_exclusion_count,
        seed=active_config.prototype_test_exclusion_seed,
    )

    dev_scored = score_prepared_split(
        dev_prepared,
        scorer,
        limit=active_config.limit_per_split,
        seed=active_config.seed,
        progress=_split_progress(progress, "score_dev"),
    )
    test_scored = score_prepared_split(
        test_prepared,
        scorer,
        limit=active_config.limit_per_split,
        seed=active_config.seed + 1,
        progress=_split_progress(progress, "score_test"),
    )

    selected = select_setting(
        dev_scored.examples,
        prior_weights=active_config.prior_weights,
        min_margins=active_config.min_margins,
    )
    dev_tuned = selected.metrics
    test_tuned = evaluate_setting(
        test_scored.examples,
        prior_weight=selected.prior_weight,
        min_margin=selected.min_margin,
    )
    dev_model_only = evaluate_setting(
        dev_scored.examples,
        prior_weight=0.0,
        min_margin=0.0,
    )
    test_model_only = evaluate_setting(
        test_scored.examples,
        prior_weight=0.0,
        min_margin=0.0,
    )
    test_outcomes = compare_examples(
        test_scored.examples,
        prior_weight=selected.prior_weight,
        min_margin=selected.min_margin,
    )
    differences = [
        int(outcome.reranked_correct) - int(outcome.baseline_correct)
        for outcome in test_outcomes
    ]
    interval = paired_bootstrap_gain_interval(
        differences,
        samples=active_config.bootstrap_samples,
        seed=active_config.seed,
    )
    report = {
        "schema_version": 1,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "LOCAL_BENCHMARK",
        "model": {
            "id": MODEL_ID,
            "revision": MODEL_REVISION,
            "task": "masked-language-model finite-candidate reranking",
            "candidate_generation": False,
        },
        "dataset": {
            "repository": UD_REPOSITORY,
            "revision": UD_REVISION,
            "license": UD_LICENSE,
            "files": {
                split: _artifact_payload(artifact) for split, artifact in artifacts.items()
            },
            "candidate_pool_size": active_config.pool_size,
            "candidate_key": "UnidicInfo reading + UPOS",
            "context_mode": "bidirectional gold UD word segmentation",
            "dev_coverage": _coverage_payload(dev_prepared),
            "test_coverage": _coverage_payload(test_prepared_all),
            "prototype_test_rows_excluded": excluded_test_rows,
            "prototype_test_exclusion_seed": active_config.prototype_test_exclusion_seed,
        },
        "sampling": {
            "limit_per_split": active_config.limit_per_split,
            "seed": active_config.seed,
            "dev_available": dev_scored.total_available,
            "dev_evaluated": len(dev_scored.examples),
            "test_available_after_exclusion": test_scored.total_available,
            "test_evaluated": len(test_scored.examples),
        },
        "selection": {
            "source": "dev only",
            "prior_weights": list(active_config.prior_weights),
            "min_margins": list(active_config.min_margins),
            "selected_prior_weight": selected.prior_weight,
            "selected_min_margin": selected.min_margin,
        },
        "metrics": {
            "dev": {
                "baseline": _baseline_payload(dev_tuned),
                "model_only": _reranked_payload(dev_model_only),
                "tuned": _reranked_payload(dev_tuned),
            },
            "sealed_test_remainder": {
                "baseline": _baseline_payload(test_tuned),
                "model_only": _reranked_payload(test_model_only),
                "tuned": _reranked_payload(test_tuned),
            },
        },
        "paired_test_statistics": {
            "absolute_gain": test_tuned.absolute_gain,
            "bootstrap_95pct_gain_interval": list(interval),
            "bootstrap_samples": active_config.bootstrap_samples,
            "bootstrap_seed": active_config.seed,
            "mcnemar_exact_p": mcnemar_exact_p(
                improved=test_tuned.improved,
                regressed=test_tuned.regressed,
            ),
            "improved": test_tuned.improved,
            "regressed": test_tuned.regressed,
        },
        "performance": {
            "device": "cpu",
            "model_load_seconds": model_load_seconds,
            "dev": _performance_payload(dev_scored),
            "test": _performance_payload(test_scored),
        },
        "environment": {
            "python": sys.version.split()[0],
            "platform": platform.platform(),
            "torch": _package_version("torch"),
            "transformers": _package_version("transformers"),
        },
        "examples": {
            "improvements": [
                _outcome_payload(outcome)
                for outcome in test_outcomes
                if not outcome.baseline_correct and outcome.reranked_correct
            ][: active_config.example_limit],
            "regressions": [
                _outcome_payload(outcome)
                for outcome in test_outcomes
                if outcome.baseline_correct and not outcome.reranked_correct
            ][: active_config.example_limit],
        },
        "claim_boundaries": [
            (
                "Accuracy is conditional on the gold surface already appearing in the "
                "top candidate pool."
            ),
            (
                "Candidate recall and romaji-to-reading conversion are not improved by "
                "this reranker."
            ),
            "UD gold bidirectional word boundaries are easier than raw incremental IME input.",
            "The baseline is train-split frequency order, not Mozc or Google Japanese Input.",
            (
                "This is local CPU evidence, not Windows TSF, device, provider, public, "
                "or human GO evidence."
            ),
        ],
    }
    return BenchmarkRun(report=report, markdown=_render_markdown(report))


def write_benchmark_outputs(
    run: BenchmarkRun,
    *,
    output_dir: Path,
    stem: str,
) -> tuple[Path, Path]:
    output_dir.mkdir(parents=True, exist_ok=True)
    json_path = output_dir / f"{stem}.json"
    markdown_path = output_dir / f"{stem}.md"
    json_path.write_text(
        json.dumps(run.report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    markdown_path.write_text(run.markdown, encoding="utf-8")
    return json_path, markdown_path


def _exclude_prototype_test_rows(
    prepared: PreparedSplit,
    *,
    count: int,
    seed: int,
) -> tuple[PreparedSplit, int]:
    actual_count = min(max(0, count), len(prepared.examples))
    indexes = list(range(len(prepared.examples)))
    random.Random(seed).shuffle(indexes)
    excluded = set(indexes[:actual_count])
    retained = tuple(
        example for index, example in enumerate(prepared.examples) if index not in excluded
    )
    return replace(prepared, examples=retained), actual_count


def _split_progress(
    progress: ProgressCallback | None,
    stage: str,
) -> Callable[[int, int], None] | None:
    if progress is None:
        return None
    return lambda completed, total: progress(stage, completed, total)


def _artifact_payload(artifact: CorpusArtifact) -> dict[str, Any]:
    return {
        "url": artifact.url,
        "sha256": artifact.sha256,
        "size_bytes": artifact.size_bytes,
        "cache_path": str(artifact.path),
    }


def _coverage_payload(prepared: PreparedSplit) -> dict[str, Any]:
    coverage = asdict(prepared.coverage)
    ambiguous = prepared.coverage.ambiguous_known_reading
    coverage["oracle_recall_at_pool"] = (
        prepared.coverage.oracle_in_pool / ambiguous if ambiguous else 0.0
    )
    return coverage


def _baseline_payload(metrics: SettingMetrics) -> dict[str, Any]:
    return {
        "correct": metrics.baseline_correct,
        "total": metrics.total,
        "accuracy": metrics.baseline_accuracy,
    }


def _reranked_payload(metrics: SettingMetrics) -> dict[str, Any]:
    return {
        "correct": metrics.reranked_correct,
        "total": metrics.total,
        "accuracy": metrics.reranked_accuracy,
        "absolute_gain_vs_baseline": metrics.absolute_gain,
        "changed": metrics.changed,
        "improved": metrics.improved,
        "regressed": metrics.regressed,
        "both_correct": metrics.both_correct,
        "both_wrong": metrics.both_wrong,
    }


def _performance_payload(scored: ScoredSplit) -> dict[str, Any]:
    latencies = sorted(example.latency_ms for example in scored.examples)
    return {
        "elapsed_seconds": scored.elapsed_seconds,
        "examples_per_second": scored.examples_per_second,
        "latency_ms_p50": _percentile(latencies, 0.50),
        "latency_ms_p95": _percentile(latencies, 0.95),
        "scoring_errors": scored.scoring_errors,
    }


def _percentile(values: list[float], fraction: float) -> float:
    if not values:
        return 0.0
    position = fraction * (len(values) - 1)
    lower = int(position)
    upper = min(lower + 1, len(values) - 1)
    weight = position - lower
    return values[lower] * (1.0 - weight) + values[upper] * weight


def _outcome_payload(outcome: ComparisonOutcome) -> dict[str, Any]:
    request = outcome.example.request
    return {
        "reading": request.reading,
        "left_context": list(request.left_context),
        "right_context": list(request.right_context),
        "candidates": [candidate.surface for candidate in request.candidates],
        "prior_scores": [candidate.prior_score for candidate in request.candidates],
        "model_scores": list(outcome.example.model_scores or ()),
        "expected": outcome.example.expected,
        "baseline": outcome.baseline_prediction,
        "reranked": outcome.reranked_prediction,
        "reason": outcome.reason,
    }


def _package_version(package: str) -> str:
    try:
        return version(package)
    except Exception:
        return "unknown"


def _render_markdown(report: dict[str, Any]) -> str:
    test = report["metrics"]["sealed_test_remainder"]
    stats = report["paired_test_statistics"]
    selection = report["selection"]
    sampling = report["sampling"]
    performance = report["performance"]["test"]
    interval = stats["bootstrap_95pct_gain_interval"]
    lines = [
        "# DeBERTa Japanese IME reranker benchmark",
        "",
        f"生成日時: {report['generated_at']}",
        "",
        "## 結果",
        "",
        "| 系 | Top-1 accuracy | 正解数 | ベースライン差 |",
        "|---|---:|---:|---:|",
        (
            f"| 頻度順ベースライン | {test['baseline']['accuracy']:.2%} | "
            f"{test['baseline']['correct']}/{test['baseline']['total']} | - |"
        ),
        (
            f"| DeBERTa 単体 | {test['model_only']['accuracy']:.2%} | "
            f"{test['model_only']['correct']}/{test['model_only']['total']} | "
            f"{test['model_only']['absolute_gain_vs_baseline']:+.2%} |"
        ),
        (
            f"| dev 調整済み混合 | {test['tuned']['accuracy']:.2%} | "
            f"{test['tuned']['correct']}/{test['tuned']['total']} | "
            f"{test['tuned']['absolute_gain_vs_baseline']:+.2%} |"
        ),
        "",
        (
            f"調整済み方式の絶対改善は {stats['absolute_gain']:+.2%} "
            f"（paired bootstrap 95% CI {interval[0]:+.2%}..{interval[1]:+.2%}、"
            f"McNemar exact p={stats['mcnemar_exact_p']:.3g}）。"
        ),
        (
            f"改善 {stats['improved']} 件、悪化 {stats['regressed']} 件。"
            f"設定は dev のみで選択し、prior weight={selection['selected_prior_weight']}、"
            f"minimum margin={selection['selected_min_margin']}。"
        ),
        "",
        "## 評価条件",
        "",
        (
            f"UD Japanese-GSD `{report['dataset']['revision']}`、候補上限 "
            f"{report['dataset']['candidate_pool_size']}。test の設計確認に使った "
            f"{report['dataset']['prototype_test_rows_excluded']} 行を除外し、"
            f"残り {sampling['test_evaluated']} 行を評価。"
        ),
        (
            f"CPU 実測は {performance['examples_per_second']:.2f} examples/s、"
            f"p50 {performance['latency_ms_p50']:.2f} ms、"
            f"p95 {performance['latency_ms_p95']:.2f} ms、"
            f"scoring errors {performance['scoring_errors']}。"
        ),
        "",
        "## 主張できないこと",
        "",
    ]
    lines.extend(f"- {boundary}" for boundary in report["claim_boundaries"])
    lines.extend(
        [
            "",
            "JSON 版には固定データハッシュ、全設定、改善/悪化例、環境情報を含む。",
            "",
        ]
    )
    return "\n".join(lines)