tests/test_benchmark.py

from __future__ import annotations

import math

import pytest

from deberta_ime import Candidate, CandidateScoringError, RerankRequest
from deberta_ime.benchmark import (
    ScoredExample,
    evaluate_setting,
    mcnemar_exact_p,
    paired_bootstrap_gain_interval,
    parse_conllu,
    prepare_examples,
    score_prepared_split,
    select_setting,
)
from deberta_ime.benchmark_runner import BenchmarkConfig, run_benchmark
from deberta_ime.benchmark_v2 import prepare_mozc_examples
from deberta_ime.mozc import MozcDictionaryIndex, build_mozc_index
from deberta_ime.ud_gsd import CorpusArtifact

CONLLU = """\
# sent_id = sample-1
# text = 川の橋を渡る。
1\t川\t川\tNOUN\t名詞\t_\t3\tnmod\t_\tSpaceAfter=No|UnidicInfo=カワ,川,川
2\tの\tの\tADP\t助詞\t_\t1\tcase\t_\tSpaceAfter=No|UnidicInfo=ノ,の,の
3\t橋\t橋\tNOUN\t名詞\t_\t4\tobj\t_\tSpaceAfter=No|UnidicInfo=ハシ,橋,橋
4\tを\tを\tADP\t助詞\t_\t3\tcase\t_\tSpaceAfter=No|UnidicInfo=ヲ,を,を
5\t渡る\t渡る\tVERB\t動詞\t_\t0\troot\t_\tSpaceAfter=No|UnidicInfo=ワタル,渡る,渡る
6\t。\t。\tPUNCT\t補助記号\t_\t5\tpunct\t_\tUnidicInfo=。,。,。

"""

TRAIN_CONLLU = """\
# sent_id = train-1
1\t箸\t箸\tNOUN\t名詞\t_\t0\troot\t_\tUnidicInfo=ハシ,箸,箸

# sent_id = train-2
1\t箸\t箸\tNOUN\t名詞\t_\t0\troot\t_\tUnidicInfo=ハシ,箸,箸

# sent_id = train-3
1\t橋\t橋\tNOUN\t名詞\t_\t0\troot\t_\tUnidicInfo=ハシ,橋,橋

"""

INFLECTED_CONLLU = """\
# sent_id = inflected-1
1\tし\t為る\tAUX\t動詞\t_\t0\troot\t_\tUnidicInfo=スル,為る,し,する,シ,,,スル,スル,する

"""


def test_parse_conllu_extracts_surface_reading_and_pos() -> None:
    sentences = parse_conllu(CONLLU)

    assert len(sentences) == 1
    assert [(token.surface, token.reading, token.upos) for token in sentences[0]] == [
        ("川", "カワ", "NOUN"),
        ("の", "ノ", "ADP"),
        ("橋", "ハシ", "NOUN"),
        ("を", "ヲ", "ADP"),
        ("渡る", "ワタル", "VERB"),
        ("。", "。", "PUNCT"),
    ]


def test_parse_conllu_uses_surface_pronunciation_instead_of_lexeme_reading() -> None:
    sentences = parse_conllu(INFLECTED_CONLLU)

    assert sentences[0][0].surface == "し"
    assert sentences[0][0].reading == "シ"


def test_prepare_examples_builds_frequency_ranked_oracle_candidate_pool() -> None:
    prepared = prepare_examples(
        parse_conllu(TRAIN_CONLLU),
        parse_conllu(CONLLU),
        pool_size=8,
    )

    assert prepared.coverage.ambiguous_known_reading == 1
    assert prepared.coverage.oracle_in_pool == 1
    assert prepared.coverage.oracle_miss == 0
    assert len(prepared.examples) == 1
    example = prepared.examples[0]
    assert example.expected == "橋"
    assert [candidate.surface for candidate in example.request.candidates] == ["箸", "橋"]
    assert [candidate.prior_score for candidate in example.request.candidates] == [
        math.log1p(2),
        math.log1p(1),
    ]
    assert example.request.left_context == ("川", "の")
    assert example.request.right_context == ("を", "渡る", "。")


def test_prepare_mozc_examples_uses_surface_reading_cost_and_context_mode(tmp_path) -> None:
    dictionary_dir = tmp_path / "dictionary_oss"
    dictionary_dir.mkdir()
    (dictionary_dir / "dictionary00.txt").write_text(
        "はし\t1\t1\t3500\t箸\nはし\t1\t1\t3800\t橋\n",
        encoding="utf-8",
    )
    index_path = tmp_path / "mozc.sqlite3"
    build_mozc_index(dictionary_dir, index_path, source_revision="fixture")

    with MozcDictionaryIndex(index_path) as index:
        prepared = prepare_mozc_examples(
            parse_conllu(CONLLU),
            index,
            pool_size=8,
            context_mode="left_only",
        )

    assert prepared.coverage.ambiguous_known_reading == 1
    assert prepared.coverage.oracle_in_pool == 1
    assert prepared.coverage.oracle_miss == 0
    example = prepared.examples[0]
    assert example.expected == "橋"
    assert [candidate.surface for candidate in example.request.candidates] == ["箸", "橋"]
    assert [candidate.prior_score for candidate in example.request.candidates] == pytest.approx(
        [0.0, -0.3]
    )
    assert example.request.left_context == ("川", "の")
    assert example.request.right_context == ()


def test_evaluate_setting_separates_improvements_and_regressions() -> None:
    def scored(
        expected: str, candidates: tuple[str, ...], scores: tuple[float, ...]
    ) -> ScoredExample:
        return ScoredExample(
            request=RerankRequest(
                reading="テスト",
                candidates=tuple(Candidate(surface) for surface in candidates),
            ),
            expected=expected,
            model_scores=scores,
            latency_ms=1.0,
        )

    examples = (
        scored("正", ("誤", "正"), (0.0, 2.0)),
        scored("正", ("正", "誤"), (0.0, 2.0)),
        scored("正", ("正", "誤"), (2.0, 0.0)),
        scored("正", ("誤一", "正", "誤二"), (0.0, 1.0, 2.0)),
    )

    metrics = evaluate_setting(examples, prior_weight=0.0, min_margin=0.0)

    assert metrics.total == 4
    assert metrics.baseline_correct == 2
    assert metrics.reranked_correct == 2
    assert metrics.improved == 1
    assert metrics.regressed == 1
    assert metrics.both_correct == 1
    assert metrics.both_wrong == 1
    assert metrics.changed == 3


def test_select_setting_uses_only_supplied_development_rows() -> None:
    request = RerankRequest(
        reading="ハシ",
        candidates=(Candidate("箸", prior_score=2.0), Candidate("橋", prior_score=0.0)),
    )
    rows = (
        ScoredExample(
            request=request,
            expected="橋",
            model_scores=(0.0, 2.0),
            latency_ms=1.0,
        ),
    )

    selected = select_setting(
        rows,
        prior_weights=(0.0, 2.0),
        min_margins=(0.0, 3.0),
    )

    assert selected.prior_weight == 0.0
    assert selected.min_margin == 0.0
    assert selected.metrics.reranked_correct == 1


def test_paired_statistics_have_known_exact_cases() -> None:
    assert mcnemar_exact_p(improved=10, regressed=0) == pytest.approx(0.001953125)
    assert mcnemar_exact_p(improved=5, regressed=5) == 1.0
    assert paired_bootstrap_gain_interval(
        [1] * 20,
        samples=200,
        seed=20260810,
    ) == (1.0, 1.0)


def test_score_prepared_split_keeps_rows_and_records_model_failures() -> None:
    prepared = prepare_examples(
        parse_conllu(TRAIN_CONLLU),
        parse_conllu(CONLLU),
        pool_size=8,
    )

    class FailingScorer:
        def score_candidates(self, request: RerankRequest) -> list[float]:
            raise RuntimeError("unavailable")

    scored = score_prepared_split(prepared, FailingScorer())

    assert len(scored.examples) == 1
    assert scored.examples[0].model_scores is None
    assert scored.scoring_errors == 1
    assert scored.total_available == 1


def test_score_prepared_split_records_stable_failure_codes() -> None:
    prepared = prepare_examples(
        parse_conllu(TRAIN_CONLLU),
        parse_conllu(CONLLU),
        pool_size=8,
    )

    class UnknownTokenScorer:
        def score_candidates(self, request: RerankRequest) -> list[float]:
            raise CandidateScoringError("unknown_token", "candidate contains unknown token")

    scored = score_prepared_split(prepared, UnknownTokenScorer())

    assert scored.examples[0].error_code == "unknown_token"


def test_benchmark_report_tunes_on_dev_and_reports_test_separately(tmp_path) -> None:
    texts = {"train": TRAIN_CONLLU, "dev": CONLLU, "test": CONLLU}

    def artifact_loader(split: str, cache_dir) -> CorpusArtifact:
        text = texts[split]
        return CorpusArtifact(
            split=split,
            path=tmp_path / f"{split}.conllu",
            url=f"https://example.invalid/{split}",
            revision="fixture",
            sha256=split * 8,
            size_bytes=len(text.encode()),
            text=text,
        )

    class ContextScorer:
        def score_candidates(self, request: RerankRequest) -> list[float]:
            return [0.0, 2.0]

    run = run_benchmark(
        ContextScorer(),
        data_dir=tmp_path,
        config=BenchmarkConfig(
            prior_weights=(0.0,),
            min_margins=(0.0,),
            bootstrap_samples=50,
            prototype_test_exclusion_count=0,
        ),
        artifact_loader=artifact_loader,
    )

    test_metrics = run.report["metrics"]["sealed_test_remainder"]
    assert run.report["selection"]["source"] == "dev only"
    assert test_metrics["baseline"]["correct"] == 0
    assert test_metrics["tuned"]["correct"] == 1
    assert run.report["paired_test_statistics"]["improved"] == 1