scripts/validate_vc02_onnx_accuracy.py

#!/usr/bin/env python3
"""Independently validate the VC02 ONNX accuracy supplement artifacts."""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import os
import pickle
import tarfile
import tempfile
from collections import Counter
from pathlib import Path
from typing import Any

import numpy as np


OUTPUTS = {
    "tflite_fp32": ("tflite_fp32_outputs.npy", np.float32),
    "tflite_public_int8": ("tflite_public_int8_outputs.npy", np.int8),
    "onnx_fp32_batch1": ("onnx_fp32_batch1_outputs.npy", np.float32),
    "onnx_fp32_optimized": ("onnx_fp32_optimized_outputs.npy", np.float32),
    "onnx_public_int8_batch1": ("onnx_public_int8_batch1_outputs.npy", np.int8),
    "onnx_public_int8_optimized": ("onnx_public_int8_optimized_outputs.npy", np.int8),
}


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def array_sha256(value: np.ndarray) -> str:
    return hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest()


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(value, handle, indent=2, sort_keys=True, allow_nan=False)
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def resolve(root: Path, value: str) -> Path:
    path = Path(value)
    return path if path.is_absolute() else root / path


def comparison(reference: np.ndarray, candidate: np.ndarray, atol: float, rtol: float, exact: bool) -> dict[str, Any]:
    if reference.shape != candidate.shape or reference.dtype != candidate.dtype:
        return {"status": "FAIL"}
    delta = candidate.astype(np.float64) - reference.astype(np.float64)
    mismatch = int(np.count_nonzero(reference != candidate))
    passed = bool(np.array_equal(reference, candidate)) if exact else bool(
        np.isfinite(reference).all() and np.isfinite(candidate).all()
        and np.allclose(candidate, reference, atol=atol, rtol=rtol)
    )
    return {
        "status": "PASS" if passed else "FAIL",
        "exact_mismatch_count": mismatch,
        "max_abs_error": float(np.abs(delta).max(initial=0.0)),
        "mean_abs_error": float(np.abs(delta).mean()) if delta.size else 0.0,
        "reference_array_sha256": array_sha256(reference),
        "candidate_array_sha256": array_sha256(candidate),
    }


def aggregate_status(
    *,
    accuracy_gate: bool,
    fp_batch_gate: bool,
    q_batch_gate: bool,
    fp_fidelity_gate: bool,
    q_fidelity_gate: bool,
    q1_reproduction_gate: bool,
) -> str:
    if (
        accuracy_gate
        and fp_batch_gate
        and q_batch_gate
        and fp_fidelity_gate
        and q_fidelity_gate
        and q1_reproduction_gate
    ):
        return "PASS"
    if accuracy_gate and fp_batch_gate and fp_fidelity_gate and q1_reproduction_gate:
        return "PARTIAL"
    return "FAIL"


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", required=True, type=Path)
    parser.add_argument("--result-dir", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    root = args.repo_root.resolve()
    result_dir = args.result_dir.resolve()
    output_path = args.output.resolve()
    summary_path = result_dir / "quality_summary.json"
    checks: list[dict[str, Any]] = []

    def check(name: str, passed: bool, detail: Any = None) -> None:
        checks.append({"name": name, "status": "PASS" if passed else "FAIL", "detail": detail})

    try:
        summary = json.loads(summary_path.read_text())
        config_path = root / "configs/evaluation/vision_classification/VC02_onnx_runtime_quality_supplement.json"
        config = json.loads(config_path.read_text())
        check("summary_model_stage", summary.get("model_id") == "VC02" and summary.get("stage") == "ONNX_RUNTIME_Q1_SUPPLEMENT")
        check("config_checksum", summary["inputs"]["supplement_config"]["sha256"] == sha256_file(config_path))
        for name, spec in config["artifacts"].items():
            path = resolve(root, spec["path"])
            check(f"artifact_{name}_sha256", path.is_file() and sha256_file(path) == spec["sha256"])
        for name, spec in config["prior_q1_evidence"].items():
            path = resolve(root, spec["path"])
            check(f"prior_q1_{name}_sha256", path.is_file() and sha256_file(path) == spec["sha256"])
        worker_package = json.loads((result_dir / "worker_results.json").read_text())
        worker_rows = worker_package.get("results", [])
        expected_jobs = ["onnx_fp32", "onnx_public_int8", "tflite_fp32", "tflite_public_int8"]
        check("worker_failures_empty", worker_package.get("failures") == [])
        check("worker_job_coverage", [row.get("job_id") for row in worker_rows] == expected_jobs)
        check("worker_summary_match", worker_rows == summary.get("worker_results"))
        expected_worker_sha = {
            "onnx_fp32": config["artifacts"]["fp32_onnx"]["sha256"],
            "onnx_public_int8": config["artifacts"]["public_int8_onnx"]["sha256"],
            "tflite_fp32": config["artifacts"]["fp32_tflite"]["sha256"],
            "tflite_public_int8": config["artifacts"]["public_int8_tflite"]["sha256"],
        }
        check(
            "worker_model_sha256",
            all(row.get("model_sha256") == expected_worker_sha.get(row.get("job_id")) for row in worker_rows),
        )

        base_path = resolve(root, config["base_quality_config"]["path"])
        base = json.loads(base_path.read_text())
        dataset_path = resolve(root, base["dataset"]["path"])
        check("dataset_sha256", sha256_file(dataset_path) == base["dataset"]["expected_sha256"])
        with tarfile.open(dataset_path, "r:gz") as archive:
            matches = [member for member in archive.getmembers() if member.isfile() and member.name.endswith("/test_batch")]
            assert len(matches) == 1
            handle = archive.extractfile(matches[0])
            assert handle is not None
            batch = pickle.loads(handle.read(), encoding="bytes")
        raw = np.asarray(batch.get(b"data", batch.get("data")), dtype=np.uint8)
        archive_labels = np.asarray(batch.get(b"labels", batch.get("labels")), dtype=np.int64)
        images_all = raw.reshape(10000, 3, 32, 32).transpose(0, 2, 3, 1)
        indices = np.asarray(np.load(resolve(root, base["protocol"]["indices_path"]), allow_pickle=False), dtype=np.int64).reshape(-1)
        expected_images = np.ascontiguousarray(images_all[indices])
        expected_labels = np.ascontiguousarray(archive_labels[indices])
        official_label_rows = []
        with resolve(root, base["protocol"]["labels_path"]).open(newline="") as handle:
            for row in csv.reader(handle):
                official_label_rows.append((row[0], int(row[1]), int(row[2])))
        expected_sample_ids = [f"cifar10_test_{int(index):05d}" for index in indices]
        images = np.load(result_dir / "semantic_inputs_uint8.npy", allow_pickle=False)
        labels = np.load(result_dir / "labels.npy", allow_pickle=False)
        check("official_input_identity", images.dtype == np.uint8 and images.shape == (200, 32, 32, 3) and np.array_equal(images, expected_images))
        check("official_label_identity", labels.dtype == np.int64 and labels.shape == (200,) and np.array_equal(labels, expected_labels))
        check("official_label_balance", Counter(labels.tolist()) == Counter({index: 20 for index in range(10)}))

        arrays: dict[str, np.ndarray] = {}
        for name, (filename, dtype) in OUTPUTS.items():
            array = np.load(result_dir / filename, allow_pickle=False)
            arrays[name] = array
            record = summary["outputs"][name]
            check(f"output_{name}_shape_dtype", array.shape == (200, 10) and array.dtype == dtype)
            check(f"output_{name}_array_sha256", array_sha256(array) == record["array_sha256"])
            check(f"output_{name}_file_sha256", sha256_file(result_dir / filename) == record["file_sha256"])

        comparisons = {
            "fp32_batch_equivalence": comparison(arrays["onnx_fp32_batch1"], arrays["onnx_fp32_optimized"], 0.0, 0.0, True),
            "public_int8_batch_equivalence": comparison(arrays["onnx_public_int8_batch1"], arrays["onnx_public_int8_optimized"], 0.0, 0.0, True),
            "fp32_tflite_to_onnx": comparison(arrays["tflite_fp32"], arrays["onnx_fp32_batch1"], 1e-4, 1e-4, False),
            "public_int8_tflite_to_onnx": comparison(arrays["tflite_public_int8"], arrays["onnx_public_int8_batch1"], 0.0, 0.0, True),
        }
        for name, observed in comparisons.items():
            claimed = summary["comparisons"][name]
            check(f"comparison_{name}_status", observed["status"] == claimed["status"], {"observed": observed, "claimed": claimed})
            check(f"comparison_{name}_mismatch", observed["exact_mismatch_count"] == claimed["exact_mismatch_count"])
            check(f"comparison_{name}_max_error", observed["max_abs_error"] == claimed["max_abs_error"])

        rows = list(csv.DictReader((result_dir / "sample_predictions.csv").open(newline="")))
        check("sample_rows_count", len(rows) == 200)
        check("sample_rows_unique_order", [int(row["sample_order"]) for row in rows] == list(range(200)))
        check("sample_rows_unique_ids", [row["sample_id"] for row in rows] == expected_sample_ids)
        check("official_filename_count", len(official_label_rows) == 200)
        for index, row in enumerate(rows):
            official_filename, official_class_count, official_label = official_label_rows[index]
            check(f"sample_{index}_filename", row["filename"] == official_filename)
            check(f"sample_{index}_class_contract", official_class_count == 10 and official_label == int(labels[index]))
            check(f"sample_{index}_label", int(row["label"]) == int(labels[index]))
            check(f"sample_{index}_input_sha", row["semantic_input_sha256"] == array_sha256(images[index]))
            for name in OUTPUTS:
                check(f"sample_{index}_{name}_prediction", int(row[f"{name}_prediction"]) == int(np.argmax(arrays[name][index])))
        predictions_record = summary["outputs"]["sample_predictions"]
        check(
            "sample_predictions_file_sha256",
            sha256_file(result_dir / "sample_predictions.csv") == predictions_record["sha256"],
        )

        threshold = float(base["metric"]["acceptance_threshold"])
        prediction_arrays = {name: np.asarray(np.argmax(array, axis=1), dtype=np.int64) for name, array in arrays.items()}
        correct_counts = {name: int(np.count_nonzero(value == labels)) for name, value in prediction_arrays.items()}
        accuracies = {name: correct_counts[name] / len(labels) for name in arrays}
        for name, value in accuracies.items():
            check(f"accuracy_{name}", value == summary["quality"][name]["accuracy"])
            check(f"correct_{name}", correct_counts[name] == summary["quality"][name]["correct"])
            check(f"sample_count_{name}", summary["quality"][name]["sample_count"] == 200)
        observed_top1 = {
            "fp32_batch_equivalence": int(np.count_nonzero(prediction_arrays["onnx_fp32_batch1"] == prediction_arrays["onnx_fp32_optimized"])),
            "public_int8_batch_equivalence": int(np.count_nonzero(prediction_arrays["onnx_public_int8_batch1"] == prediction_arrays["onnx_public_int8_optimized"])),
            "fp32_tflite_to_onnx": int(np.count_nonzero(prediction_arrays["tflite_fp32"] == prediction_arrays["onnx_fp32_batch1"])),
            "public_int8_tflite_to_onnx": int(np.count_nonzero(prediction_arrays["tflite_public_int8"] == prediction_arrays["onnx_public_int8_batch1"])),
        }
        for name, count in observed_top1.items():
            check(f"top1_{name}", count == summary["comparisons"][name]["top1_agreement_count"])

        prior_rows = list(csv.DictReader(resolve(root, config["prior_q1_evidence"]["sample_predictions"]["path"]).open(newline="")))
        prior_by_id = {row["sample_id"]: row for row in prior_rows}
        check("prior_q1_rows_unique", len(prior_rows) == len(prior_by_id) == 200)
        prior_fp_agreement = sum(
            int(prior_by_id[sample_id]["fp32_prediction"]) == int(prediction_arrays["tflite_fp32"][index])
            for index, sample_id in enumerate(expected_sample_ids)
        )
        prior_q_agreement = sum(
            int(prior_by_id[sample_id]["public_int8_prediction"]) == int(prediction_arrays["tflite_public_int8"][index])
            for index, sample_id in enumerate(expected_sample_ids)
        )
        q1_reproduction_pass = prior_fp_agreement == prior_q_agreement == 200
        check("q1_reproduction_status", summary["q1_reproduction"]["status"] == ("PASS" if q1_reproduction_pass else "FAIL"))
        check("q1_reproduction_fp32_count", summary["q1_reproduction"]["fp32_prediction_agreement_count"] == prior_fp_agreement)
        check("q1_reproduction_quant_count", summary["q1_reproduction"]["public_int8_prediction_agreement_count"] == prior_q_agreement)
        official_q1_pass = accuracies["tflite_public_int8"] >= threshold
        onnx_accuracy_pass = accuracies["onnx_fp32_batch1"] >= threshold and accuracies["onnx_public_int8_batch1"] >= threshold
        fp_batch_pass = comparisons["fp32_batch_equivalence"]["status"] == "PASS"
        q_batch_pass = comparisons["public_int8_batch_equivalence"]["status"] == "PASS"
        fp_fidelity_pass = comparisons["fp32_tflite_to_onnx"]["status"] == "PASS" and observed_top1["fp32_tflite_to_onnx"] == 200
        q_fidelity_pass = comparisons["public_int8_tflite_to_onnx"]["status"] == "PASS" and observed_top1["public_int8_tflite_to_onnx"] == 200
        expected_status = aggregate_status(
            accuracy_gate=official_q1_pass and onnx_accuracy_pass,
            fp_batch_gate=fp_batch_pass,
            q_batch_gate=q_batch_pass,
            fp_fidelity_gate=fp_fidelity_pass,
            q_fidelity_gate=q_fidelity_pass,
            q1_reproduction_gate=q1_reproduction_pass,
        )
        check("official_q1_status", summary["official_q1_acceptance_status"] == ("PASS" if official_q1_pass else "FAIL"))
        check("onnx_accuracy_status", summary["onnx_supplement_accuracy_status"] == ("PASS" if onnx_accuracy_pass else "FAIL"))
        check("overall_status", summary["status"] == expected_status)
        check("conversion_fidelity_status", summary["conversion_fidelity_status"] == ("PASS" if fp_fidelity_pass and q_fidelity_pass else "FAIL"))
        check("optimized_execution_status", summary["optimized_execution_status"] == ("PASS" if fp_batch_pass and q_batch_pass else "FAIL"))
        check("optimized_fp32_accepted", summary["optimized_outputs_accepted"]["fp32"] is fp_batch_pass)
        check("optimized_quant_accepted", summary["optimized_outputs_accepted"]["public_int8"] is q_batch_pass)
        check("policy_no_latency", summary["policy"]["latency_benchmark"] is False)
        check("policy_no_converter", summary["policy"]["converter_run"] is False)
        check("policy_no_codegen", summary["policy"]["mlir_lowering_or_codegen_run"] is False)
    except Exception as error:
        check("validator_exception", False, f"{type(error).__name__}: {error}")

    failed = [row for row in checks if row["status"] != "PASS"]
    report = {
        "schema_version": "1.0",
        "model_id": "VC02",
        "stage": "ONNX_RUNTIME_Q1_SUPPLEMENT_INDEPENDENT_VALIDATION",
        "status": "PASS" if not failed else "FAIL",
        "checks_total": len(checks),
        "checks_passed": len(checks) - len(failed),
        "checks_failed": len(failed),
        "checks": checks,
        "validator_independence": {
            "model_runtime_invoked": False,
            "evaluator_module_imported": False,
            "outputs_reloaded_and_recomputed": True,
        },
    }
    atomic_json(output_path, report)
    print(json.dumps({key: report[key] for key in ("status", "checks_total", "checks_failed")}, sort_keys=True))
    return 0 if report["status"] == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())