tests/test_lm04_raid_quality.py

from __future__ import annotations

import json
from pathlib import Path

import pytest

from scripts.stages.evaluate_lm04_raid_quality import (
    BertWordPieceTokenizer,
    auc_score,
    deterministic_rank,
    largest_remainder_allocations,
    machine_score,
    official_find_threshold,
)


ROOT = Path(__file__).resolve().parents[1]
VOCAB = ROOT / "models/language_modeling_text_classification/LM04/config/vocab.txt"


@pytest.mark.skipif(not VOCAB.is_file(), reason="compact repository omits generated LM04 tokenizer files")
def test_fixed_prompt_matches_distributed_vocab_ids() -> None:
    tokenizer = BertWordPieceTokenizer(VOCAB)
    assert tokenizer.encode("this is a story") == [101, 2023, 2003, 1037, 2466, 102]


@pytest.mark.skipif(not VOCAB.is_file(), reason="compact repository omits generated LM04 tokenizer files")
def test_uncased_accent_wordpiece_and_right_truncation() -> None:
    tokenizer = BertWordPieceTokenizer(VOCAB)
    assert tokenizer.basic_tokenize("Héllo, Café!") == ["hello", ",", "cafe", "!"]
    ids = tokenizer.encode(" ".join(f"token{i % 97}" for i in range(900)))
    assert len(ids) == 512
    assert ids[0] == 101 and ids[-1] == 102


def test_machine_score_inverts_human_logit_direction() -> None:
    assert machine_score(-2.0) > machine_score(0.0) > machine_score(2.0)


def test_auc_handles_ties_by_average_rank() -> None:
    assert auc_score([0, 0, 1, 1], [0.1, 0.2, 0.8, 0.9]) == 1.0
    assert auc_score([0, 1], [0.5, 0.5]) == 0.5


def test_largest_remainder_allocation_is_exact_and_deterministic() -> None:
    counts = {("code", "a"): 3, ("german", "b"): 7}
    assert largest_remainder_allocations(counts, 5) == {("code", "a"): 2, ("german", "b"): 3}
    assert deterministic_rank("p", "x") == deterministic_rank("p", "x")
    assert deterministic_rank("p", "x") != deterministic_rank("p", "y")


def test_official_threshold_search_meets_or_selects_closest_fpr() -> None:
    human = [index / 1000 for index in range(1000)]
    threshold, observed = official_find_threshold(human, 0.05, 0.0005)
    assert threshold >= 0.0
    assert abs(observed - 0.05) <= 0.001


@pytest.mark.skipif(
    not (ROOT / "results/lm04_raid_extra_ood/tokenizer_reference_equivalence.json").is_file(),
    reason="compact repository omits materialized LM04 evaluation results",
)
def test_pinned_bert_tokenizer_fast_reference_equivalence() -> None:
    path = ROOT / "results/lm04_raid_extra_ood/tokenizer_reference_equivalence.json"
    document = json.loads(path.read_text(encoding="utf-8"))
    assert document["status"] == "PASS"
    assert document["reference"]["transformers_version"] == "4.57.6"
    assert document["reference"]["tokenizers_version"] == "0.22.2"
    assert document["case_count"] >= 14
    assert document["all_input_tensors_exact"] is True
    assert all(item["status"] == "PASS" and item["mismatch_count"] == 0
               for item in document["comparisons"].values())
    assert any(case["case_name"].startswith("dataset:czech") for case in document["cases"])
    assert any(case["case_name"].startswith("dataset:german") for case in document["cases"])
    assert any(case["reaches_max_length"] for case in document["cases"])


def test_actual_lm04_quality_package_when_present() -> None:
    validation_path = ROOT / "results/lm04_raid_extra_ood/evaluation_v1/validation.json"
    if not validation_path.is_file():
        import pytest
        pytest.skip("LM04 RAID-extra package has not been generated")
    validation = json.loads(validation_path.read_text(encoding="utf-8"))
    assert validation["status"] == "PASS"
    assert validation["checks_failed"] == 0