code/tests/test_acl_arr_supplement.py

from __future__ import annotations

import csv
import json
from pathlib import Path
from statistics import mean

from scripts.export_acl_arr_supplement import export


REPO = Path(__file__).resolve().parents[1]


def _rows(path: Path) -> list[dict[str, str]]:
    with path.open(newline="") as handle:
        return list(csv.DictReader(handle))


def test_export_contains_complete_existing_grids(tmp_path: Path) -> None:
    broad_path, alpha_path, aggregate_path, manifest_path = export(REPO, tmp_path)
    broad, alpha = _rows(broad_path), _rows(alpha_path)
    aggregates = json.loads(aggregate_path.read_text())
    manifest = _rows(manifest_path)

    assert len(broad) == 1020
    assert len(alpha) == 560
    assert {row["status"] for row in broad + alpha} == {"complete"}

    source_cache: dict[str, dict[str, object]] = {}
    for row in broad + alpha:
        source = row["source"]
        source_cache.setdefault(source, json.loads((REPO / source).read_text()))
        key = f"{row['method']}__rank{row['rank']}__seed{row['seed']}"
        assert float(source_cache[source][key]["eval_accuracy"]) == float(row["accuracy"])

    input_tail = [row for row in broad if row["method"] == "fpeft_low"]
    peft = [row for row in broad if row["method"] == "peft_default"]
    assert len(input_tail) == len(peft) == 210
    assert sum(float(row["accuracy"]) < 0.70 for row in input_tail) == 11
    assert sum(float(row["accuracy"]) < 0.70 for row in peft) == 29

    peft_by_block = {
        block: mean(float(row["accuracy"]) for row in peft if row["block"] == block)
        for block in {row["block"] for row in peft}
    }
    input_by_block = {
        block: mean(float(row["accuracy"]) for row in input_tail if row["block"] == block)
        for block in {row["block"] for row in input_tail}
    }
    block_delta = mean(
        100 * (input_by_block[block] - peft_by_block[block])
        for block in peft_by_block
    )
    assert round(block_delta, 2) == 2.23

    assert aggregates["selection_audit"]["all_per_block_oracle"]["calibration_low_cells"] == 7
    assert round(
        aggregates["matched_scale"]["fpeft_low__vs__random_orthogonal"]["mean_delta_pp"],
        2,
    ) == -0.03
    assert [
        round(value, 2)
        for value in aggregates["matched_scale"][
            "fpeft_low__vs__random_orthogonal"
        ]["block_bootstrap_95_pp"]
    ] == [-0.47, 0.35]
    assert round(
        aggregates["extreme_multiplier_interactions"][
            "fpeft_bi_high__vs__peft_scaled_random"
        ]["mean_difference_in_differences_pp"],
        2,
    ) == -0.99
    assert [
        round(value, 2)
        for value in aggregates["extreme_multiplier_interactions"][
            "fpeft_bi_high__vs__peft_scaled_random"
        ]["block_bootstrap_95_pp"]
    ] == [-2.88, 0.30]

    assert len(manifest) == 432
    assert {row["dataset_revision"] for row in manifest} == {"unknown_not_archived"}
    assert {row["factor_checkpoint_revision"] for row in manifest} == {"unknown_not_archived"}
    assert {int(row["count_a"]) for row in manifest if row["model"] == "llama"} == {53248}