src/deberta_ime/benchmark_v2_runner.py

from __future__ import annotations

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

from .benchmark import (
    ComparisonOutcome,
    ScoredSplit,
    SettingMetrics,
    compare_examples,
    evaluate_setting,
    mcnemar_exact_p,
    paired_bootstrap_gain_interval,
    parse_conllu,
    score_prepared_split,
    select_setting,
)
from .benchmark_v2 import prepare_mozc_examples
from .deberta import MODEL_ID, MODEL_REVISION
from .domain import CandidateScorer
from .mozc import MozcDictionaryIndex
from .ud_gsd import UD_LICENSE, UD_REPOSITORY, CorpusArtifact


@dataclass(frozen=True, slots=True)
class FrozenProfile:
    name: str
    context_mode: str
    prior_weight: float
    min_margin: float
    selected_on: str


@dataclass(frozen=True, slots=True)
class V2BenchmarkConfig:
    pool_size: int = 8
    limit: int | None = None
    seed: int = 20260810
    prior_weights: tuple[float, ...] = (
        0.0,
        0.25,
        0.5,
        0.75,
        1.0,
        1.5,
        2.0,
        3.0,
        4.0,
        6.0,
        8.0,
    )
    min_margins: tuple[float, ...] = (0.1, 0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0)
    bootstrap_samples: int = 5000
    example_limit: int = 20


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


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


def write_v2_outputs(
    run: V2BenchmarkRun,
    *,
    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 run_dev_selection(
    scorer: CandidateScorer,
    *,
    index: MozcDictionaryIndex,
    artifact: CorpusArtifact,
    context_mode: str,
    config: V2BenchmarkConfig | None = None,
    model_load_seconds: float | None = None,
    progress: ProgressCallback | None = None,
) -> V2BenchmarkRun:
    active_config = config or V2BenchmarkConfig()
    prepared = prepare_mozc_examples(
        parse_conllu(artifact.text),
        index,
        pool_size=active_config.pool_size,
        context_mode=context_mode,
    )
    scored = score_prepared_split(
        prepared,
        scorer,
        limit=active_config.limit,
        seed=active_config.seed,
        progress=progress,
    )
    selected = select_setting(
        scored.examples,
        prior_weights=active_config.prior_weights,
        min_margins=active_config.min_margins,
    )
    model_only = evaluate_setting(
        scored.examples,
        prior_weight=0.0,
        min_margin=0.0,
    )
    ambiguous = prepared.coverage.ambiguous_known_reading
    report: dict[str, Any] = {
        "schema_version": 2,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "LOCAL_BENCHMARK",
        "stage": "DEV_SELECTION",
        "model": {
            "id": MODEL_ID,
            "revision": MODEL_REVISION,
            "task": "masked-language-model finite-candidate reranking",
            "candidate_generation": False,
        },
        "candidate_source": {
            "kind": "Mozc OSS dictionary SQLite index",
            "manifest": asdict(index.manifest),
            "pool_size": active_config.pool_size,
            "prior_score": "-(cost - best_cost) / 1000",
        },
        "dataset": {
            "repository": UD_REPOSITORY,
            "revision": artifact.revision,
            "license": UD_LICENSE,
            "file": {
                "url": artifact.url,
                "sha256": artifact.sha256,
                "size_bytes": artifact.size_bytes,
                "cache_path": str(artifact.path),
            },
        },
        "selection": {
            "source": "GSD dev only",
            "dataset_revision": artifact.revision,
            "context_mode": context_mode,
            "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,
        },
        "sampling": {
            "limit": active_config.limit,
            "seed": active_config.seed,
            "available": scored.total_available,
            "evaluated": len(scored.examples),
        },
        "coverage": _coverage_payload(prepared.coverage),
        "conditional_metrics": {
            "baseline": _baseline_payload(selected.metrics),
            "model_only": _reranked_payload(model_only),
            "selected": _reranked_payload(selected.metrics),
        },
        "end_to_end_ambiguous": {
            "denominator": ambiguous,
            "baseline_accuracy": (
                selected.metrics.baseline_correct / ambiguous if ambiguous else 0.0
            ),
            "model_only_accuracy": (
                model_only.reranked_correct / ambiguous if ambiguous else 0.0
            ),
            "selected_accuracy": (
                selected.metrics.reranked_correct / ambiguous if ambiguous else 0.0
            ),
        },
        "performance": _performance_payload(scored, model_load_seconds=model_load_seconds),
        "environment": _environment_payload(),
        "claim_boundaries": [
            "This report selects hyperparameters and is not external evaluation evidence.",
            "Conditional accuracy requires the expected surface in Mozc top-k.",
            "PUD and AJIMEE are not read or scored by this function.",
        ],
    }
    return V2BenchmarkRun(report=report, markdown=_render_dev_markdown(report))


def run_external_evaluation(
    scorer: CandidateScorer,
    *,
    index: MozcDictionaryIndex,
    artifact: CorpusArtifact,
    dataset_repository: str,
    dataset_license: str,
    profile: FrozenProfile,
    config: V2BenchmarkConfig | None = None,
    model_load_seconds: float | None = None,
    progress: ProgressCallback | None = None,
) -> V2BenchmarkRun:
    active_config = config or V2BenchmarkConfig()
    prepared = prepare_mozc_examples(
        parse_conllu(artifact.text),
        index,
        pool_size=active_config.pool_size,
        context_mode=profile.context_mode,
    )
    scored = score_prepared_split(
        prepared,
        scorer,
        limit=active_config.limit,
        seed=active_config.seed,
        progress=progress,
    )
    tuned = evaluate_setting(
        scored.examples,
        prior_weight=profile.prior_weight,
        min_margin=profile.min_margin,
    )
    model_only = evaluate_setting(
        scored.examples,
        prior_weight=0.0,
        min_margin=0.0,
    )
    outcomes = compare_examples(
        scored.examples,
        prior_weight=profile.prior_weight,
        min_margin=profile.min_margin,
    )
    differences = [
        int(outcome.reranked_correct) - int(outcome.baseline_correct) for outcome in outcomes
    ]
    interval = paired_bootstrap_gain_interval(
        differences,
        samples=active_config.bootstrap_samples,
        seed=active_config.seed,
    )
    ambiguous = prepared.coverage.ambiguous_known_reading
    report: dict[str, Any] = {
        "schema_version": 2,
        "generated_at": datetime.now(UTC).isoformat(),
        "status": "LOCAL_BENCHMARK",
        "stage": "EXTERNAL_EVALUATION",
        "model": {
            "id": MODEL_ID,
            "revision": MODEL_REVISION,
            "task": "masked-language-model finite-candidate reranking",
            "candidate_generation": False,
        },
        "candidate_source": {
            "kind": "Mozc OSS dictionary SQLite index",
            "manifest": asdict(index.manifest),
            "pool_size": active_config.pool_size,
            "prior_score": "-(cost - best_cost) / 1000",
        },
        "dataset": {
            "repository": dataset_repository,
            "revision": artifact.revision,
            "license": dataset_license,
            "file": {
                "url": artifact.url,
                "sha256": artifact.sha256,
                "size_bytes": artifact.size_bytes,
                "cache_path": str(artifact.path),
            },
        },
        "profile": {**asdict(profile), "source": "frozen"},
        "sampling": {
            "limit": active_config.limit,
            "seed": active_config.seed,
            "available": scored.total_available,
            "evaluated": len(scored.examples),
        },
        "coverage": _coverage_payload(prepared.coverage),
        "conditional_metrics": {
            "baseline": _baseline_payload(tuned),
            "model_only": _reranked_payload(model_only),
            "tuned": _reranked_payload(tuned),
        },
        "end_to_end_ambiguous": {
            "denominator": ambiguous,
            "baseline_accuracy": tuned.baseline_correct / ambiguous if ambiguous else 0.0,
            "model_only_accuracy": (
                model_only.reranked_correct / ambiguous if ambiguous else 0.0
            ),
            "tuned_accuracy": tuned.reranked_correct / ambiguous if ambiguous else 0.0,
        },
        "paired_statistics": {
            "absolute_gain": 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=tuned.improved,
                regressed=tuned.regressed,
            ),
            "improved": tuned.improved,
            "regressed": tuned.regressed,
        },
        "performance": _performance_payload(scored, model_load_seconds=model_load_seconds),
        "environment": _environment_payload(),
        "examples": _examples_payload(outcomes, limit=active_config.example_limit),
        "claim_boundaries": [
            "Conditional accuracy requires the expected surface in Mozc top-k.",
            "Raw romaji conversion and sentence-level candidate generation are not evaluated.",
            "Gold UD word boundaries are supplied as context.",
            "This is local CPU evidence, not TSF, device, public, or Human GO evidence.",
        ],
    }
    return V2BenchmarkRun(report=report, markdown=_render_external_markdown(report))


def _coverage_payload(coverage: Any) -> dict[str, Any]:
    ambiguous = coverage.ambiguous_known_reading
    return {
        "eligible": coverage.eligible_tokens,
        "ambiguous": ambiguous,
        "oracle_in_pool": coverage.oracle_in_pool,
        "oracle_miss": coverage.oracle_miss,
        "candidate_recall_at_k": coverage.oracle_in_pool / ambiguous if ambiguous else 0.0,
    }


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,
    *,
    model_load_seconds: float | None,
) -> dict[str, Any]:
    latencies = sorted(example.latency_ms for example in scored.examples)
    failure_counts = Counter(
        example.error_code for example in scored.examples if example.error_code is not None
    )
    return {
        "device": "cpu",
        "model_load_seconds": model_load_seconds,
        "elapsed_seconds": scored.elapsed_seconds,
        "examples_per_second": scored.examples_per_second,
        "latency_ms_p50": _percentile(latencies, 0.5),
        "latency_ms_p95": _percentile(latencies, 0.95),
        "scoring_errors": scored.scoring_errors,
        "scoring_failure_counts": dict(sorted(failure_counts.items())),
        "scoring_error_rate": (
            scored.scoring_errors / len(scored.examples) if scored.examples else 0.0
        ),
    }


def _environment_payload() -> dict[str, str]:
    return {
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "torch": _package_version("torch"),
        "transformers": _package_version("transformers"),
    }


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


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 _examples_payload(
    outcomes: tuple[ComparisonOutcome, ...],
    *,
    limit: int,
) -> dict[str, Any]:
    def payload(outcome: ComparisonOutcome) -> dict[str, Any]:
        return {
            "reading": outcome.example.request.reading,
            "expected": outcome.example.expected,
            "baseline": outcome.baseline_prediction,
            "reranked": outcome.reranked_prediction,
            "reason": outcome.reason,
        }

    return {
        "improvements": [
            payload(outcome)
            for outcome in outcomes
            if not outcome.baseline_correct and outcome.reranked_correct
        ][:limit],
        "regressions": [
            payload(outcome)
            for outcome in outcomes
            if outcome.baseline_correct and not outcome.reranked_correct
        ][:limit],
    }


def _render_external_markdown(report: dict[str, Any]) -> str:
    metrics = report["conditional_metrics"]
    coverage = report["coverage"]
    profile = report["profile"]
    end_to_end = report["end_to_end_ambiguous"]
    statistics_payload = report["paired_statistics"]
    performance = report["performance"]
    interval = statistics_payload["bootstrap_95pct_gain_interval"]
    failures = json.dumps(
        performance["scoring_failure_counts"], ensure_ascii=False, sort_keys=True
    )
    lines = [
        "# DeBERTa Japanese IME v0.2 external evaluation",
        "",
        f"生成日時: {report['generated_at']}",
        "",
        (
            f"Profile: `{profile['name']}` ({profile['context_mode']}), "
            f"prior weight={profile['prior_weight']}, min margin={profile['min_margin']}。"
        ),
        f"選択元: `{profile['selected_on']}`。外部データで再調整していない。",
        "",
        "## 候補coverage",
        "",
        (
            f"eligible {coverage['eligible']}、複数候補 {coverage['ambiguous']}、"
            f"Mozc top-{report['candidate_source']['pool_size']} candidate recall "
            f"{coverage['candidate_recall_at_k']:.2%} "
            f"({coverage['oracle_in_pool']}/{coverage['ambiguous']})、"
            f"candidate miss {coverage['oracle_miss']}。"
        ),
        "",
        "## Oracle-in-pool 条件付きTop-1",
        "",
        "| 系 | Accuracy | 正解数 | ベースライン差 |",
        "|---|---:|---:|---:|",
        (
            f"| Mozc cost | {metrics['baseline']['accuracy']:.2%} | "
            f"{metrics['baseline']['correct']}/{metrics['baseline']['total']} | - |"
        ),
        (
            f"| DeBERTa単体 | {metrics['model_only']['accuracy']:.2%} | "
            f"{metrics['model_only']['correct']}/{metrics['model_only']['total']} | "
            f"{metrics['model_only']['absolute_gain_vs_baseline']:+.2%} |"
        ),
        (
            f"| 凍結済み混合 | {metrics['tuned']['accuracy']:.2%} | "
            f"{metrics['tuned']['correct']}/{metrics['tuned']['total']} | "
            f"{metrics['tuned']['absolute_gain_vs_baseline']:+.2%} |"
        ),
        "",
        (
            f"改善 {statistics_payload['improved']}件、悪化 "
            f"{statistics_payload['regressed']}件。paired bootstrap 95% CI "
            f"{interval[0]:+.2%}..{interval[1]:+.2%}、"
            f"McNemar exact p={statistics_payload['mcnemar_exact_p']:.3g}。"
        ),
        "",
        "## Candidate miss込みTop-1",
        "",
        (
            f"複数候補全体 {end_to_end['denominator']}件で、Mozc "
            f"{end_to_end['baseline_accuracy']:.2%}、凍結済み混合 "
            f"{end_to_end['tuned_accuracy']:.2%}。candidate miss は誤りとして数えた。"
        ),
        "",
        "## CPUとfail-closed",
        "",
        (
            f"{performance['examples_per_second']:.2f} examples/s、p50 "
            f"{performance['latency_ms_p50']:.2f} ms、p95 "
            f"{performance['latency_ms_p95']:.2f} ms。scoring errors "
            f"{performance['scoring_errors']} ({performance['scoring_error_rate']:.2%})、"
            f"内訳 `{failures}`。すべてMozc順位維持として分母に含めた。"
        ),
        "",
        "## 主張できないこと",
        "",
    ]
    lines.extend(f"- {boundary}" for boundary in report["claim_boundaries"])
    lines.append("")
    return "\n".join(lines)


def _render_dev_markdown(report: dict[str, Any]) -> str:
    metrics = report["conditional_metrics"]
    coverage = report["coverage"]
    selection = report["selection"]
    return "\n".join(
        [
            "# DeBERTa Japanese IME v0.2 dev selection",
            "",
            f"生成日時: {report['generated_at']}",
            "",
            f"Context: `{selection['context_mode']}`",
            (
                f"Selected: prior_weight={selection['selected_prior_weight']}, "
                f"min_margin={selection['selected_min_margin']}"
            ),
            "",
            (
                f"Mozc top-{report['candidate_source']['pool_size']} candidate recall: "
                f"{coverage['candidate_recall_at_k']:.2%} "
                f"({coverage['oracle_in_pool']}/{coverage['ambiguous']})"
            ),
            "",
            "| 系 | Oracle-in-pool Top-1 | ベースライン差 |",
            "|---|---:|---:|",
            f"| Mozc cost | {metrics['baseline']['accuracy']:.2%} | - |",
            (
                f"| DeBERTa単体 | {metrics['model_only']['accuracy']:.2%} | "
                f"{metrics['model_only']['absolute_gain_vs_baseline']:+.2%} |"
            ),
            (
                f"| dev選択混合 | {metrics['selected']['accuracy']:.2%} | "
                f"{metrics['selected']['absolute_gain_vs_baseline']:+.2%} |"
            ),
            "",
            "DEV_SELECTION のため、外部評価の主張には使わない。",
            "",
        ]
    )