scripts/audit_math_ink_06_online_error_consensus.py

"""세 seed의 online 오류 합의도를 계산해 데이터·모델·seed 분산 병목을 분리한다."""

from __future__ import annotations

import argparse
from collections import Counter, defaultdict
from datetime import datetime, timezone
import json
from pathlib import Path
import sys
from typing import Sequence

import torch

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.trajectory_sequence import shape_family, visual_label_family
from scripts.train_math_ink_06_p_boundary_auxiliary import _load_encoder06
from scripts.train_math_ink_06_skeleton_adapter import _resolve_device06


def _parse_args() -> argparse.Namespace:
    """필요 변수: feature cache·seed별 checkpoint. 작동 원리: 동일 분모의 합의 오류 감사 CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Audit Math Ink 0.6 online error consensus")
    parser.add_argument("--feature-cache", type=Path, required=True)
    parser.add_argument("--base-checkpoint", type=Path, action="append", required=True)
    parser.add_argument("--adapter-checkpoint", type=Path, action="append", required=True)
    parser.add_argument("--source-counts", type=Path)
    parser.add_argument("--source-split", choices=("validation", "paired_test"), default="validation")
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if len(args.base_checkpoint) != len(args.adapter_checkpoint):
        raise ValueError("base와 adapter checkpoint 개수는 같아야 합니다.")
    if len(args.base_checkpoint) < 2:
        raise ValueError("seed 합의 감사에는 checkpoint 두 개 이상이 필요합니다.")
    return args


def _source_vector06(
    samples: int, source_counts_path: Path | None, split: str,
) -> list[str]:
    """필요 변수: 총 표본 수·source별 개수 보고서. 작동 원리: cache 생성 순서와 같은 source 구간을 복원한다."""

    if source_counts_path is None:
        return ["unknown"] * samples
    payload = json.loads(source_counts_path.read_text(encoding="utf-8"))
    counts = payload.get("paired_source_counts", payload)
    values: list[str] = []
    for source in ("hwrt", "uci-uji-pen-v1", "uci-uji-pen-v2"):
        if source not in counts:
            continue
        count = int(counts[source][split])
        values.extend([source] * count)
    if len(values) != samples:
        raise ValueError(f"source count 합계가 cache와 다릅니다: {len(values)} != {samples}")
    return values


def summarize_online_consensus06(
    logits_by_seed: Sequence[torch.Tensor],
    targets: torch.Tensor,
    writers: torch.Tensor,
    labels: Sequence[str],
    sources: Sequence[str],
) -> dict:
    """필요 변수: seed별 logit·정답·writer/source. 작동 원리: 공통 오류와 ensemble 상한을 한 분모에서 계산한다."""

    if not logits_by_seed:
        raise ValueError("seed logit이 비어 있습니다.")
    samples = len(targets)
    if any(len(logits) != samples for logits in logits_by_seed):
        raise ValueError("seed logit 표본 수가 서로 다릅니다.")
    if len(writers) != samples or len(sources) != samples:
        raise ValueError("writer/source 표본 수가 target과 다릅니다.")

    predictions = torch.stack([logits.argmax(dim=-1) for logits in logits_by_seed])
    correctness = predictions.eq(targets.unsqueeze(0))
    ensemble_logits = torch.stack(
        [logits.log_softmax(dim=-1) for logits in logits_by_seed],
    ).logsumexp(dim=0)
    ensemble_prediction = ensemble_logits.argmax(dim=-1)
    ensemble_top5 = ensemble_logits.topk(min(5, ensemble_logits.shape[-1]), dim=-1).indices
    ensemble_correct = ensemble_prediction.eq(targets)
    oracle_correct = correctness.any(dim=0)
    unanimous_wrong = correctness.logical_not().all(dim=0)
    unanimous_same_prediction = predictions.eq(predictions[:1]).all(dim=0)
    shape_families = tuple(shape_family(str(label)) for label in labels)
    visual_families = tuple(visual_label_family(str(label)) for label in labels)
    shape_correct = torch.tensor([
        shape_families[truth] == shape_families[predicted]
        for truth, predicted in zip(targets.tolist(), ensemble_prediction.tolist(), strict=True)
    ])
    visual_correct = torch.tensor([
        visual_families[truth] == visual_families[predicted]
        for truth, predicted in zip(targets.tolist(), ensemble_prediction.tolist(), strict=True)
    ])
    exact_wrong = ensemble_correct.logical_not()

    per_label: dict[str, dict] = {}
    confusion = Counter()
    for index, label in enumerate(labels):
        mask = targets.eq(index)
        count = int(mask.sum())
        if count == 0:
            continue
        correct = int(ensemble_correct[mask].sum())
        oracle = int(oracle_correct[mask].sum())
        common = int(unanimous_wrong[mask].sum())
        per_label[str(label)] = {
            "samples": count,
            "ensemble_top1": correct / count,
            "seed_oracle_top1": oracle / count,
            "unanimous_wrong_rate": common / count,
        }
    for truth, predicted in zip(targets.tolist(), ensemble_prediction.tolist(), strict=True):
        if truth != predicted:
            confusion[(str(labels[truth]), str(labels[predicted]))] += 1

    def _slice_rows(keys: Sequence[str | int]) -> dict[str, dict]:
        """필요 변수: source 또는 writer key. 작동 원리: 동일 ensemble 지표를 slice별로 집계한다."""

        grouped: dict[str, list[int]] = defaultdict(list)
        for index, key in enumerate(keys):
            grouped[str(key)].append(index)
        rows = {}
        for key, indices in grouped.items():
            mask = torch.tensor(indices, dtype=torch.long)
            count = len(indices)
            rows[key] = {
                "samples": count,
                "ensemble_top1": float(ensemble_correct[mask].float().mean()),
                "seed_oracle_top1": float(oracle_correct[mask].float().mean()),
                "unanimous_wrong_rate": float(unanimous_wrong[mask].float().mean()),
            }
        return rows

    per_writer = _slice_rows(writers.tolist())
    eligible_writers = [
        row["ensemble_top1"] for row in per_writer.values() if int(row["samples"]) >= 10
    ]
    weakest_labels = sorted(
        (
            {"label": label, **row}
            for label, row in per_label.items() if int(row["samples"]) >= 5
        ),
        key=lambda row: (float(row["ensemble_top1"]), -int(row["samples"]), str(row["label"])),
    )[:30]
    top_confusions = [
        {"truth": truth, "predicted": predicted, "count": count}
        for (truth, predicted), count in confusion.most_common(40)
    ]
    return {
        "samples": samples,
        "seed_count": len(logits_by_seed),
        "seed_top1": [
            float(correct.float().mean()) for correct in correctness
        ],
        "ensemble_top1": float(ensemble_correct.float().mean()),
        "ensemble_top5": float(
            ensemble_top5.eq(targets[:, None]).any(dim=-1).float().mean()
        ),
        "shape_family_top1": float(shape_correct.float().mean()),
        "visual_family_top1": float(visual_correct.float().mean()),
        "exact_error_visual_family_recoverable_rate": float(
            visual_correct[exact_wrong].float().mean() if exact_wrong.any() else 0.0
        ),
        "exact_error_visual_family_recoverable_pp": float(
            (visual_correct & exact_wrong).float().mean() * 100.0
        ),
        "seed_oracle_top1": float(oracle_correct.float().mean()),
        "all_seed_wrong_rate": float(unanimous_wrong.float().mean()),
        "all_seed_same_wrong_rate": float(
            (unanimous_wrong & unanimous_same_prediction).float().mean()
        ),
        "recoverable_by_seed_choice_pp": float(
            (oracle_correct.float().mean() - ensemble_correct.float().mean()) * 100.0
        ),
        "eligible_writer_count": len(eligible_writers),
        "eligible_writer_floor": min(eligible_writers, default=0.0),
        "per_source": _slice_rows(sources),
        "weakest_labels_min5": weakest_labels,
        "top_confusions": top_confusions,
    }


def _infer_logits06(
    base_paths: Sequence[Path],
    adapter_paths: Sequence[Path],
    features: torch.Tensor,
    *,
    device: torch.device,
    batch_size: int,
) -> tuple[list[torch.Tensor], tuple[str, ...], tuple[str, ...]]:
    """필요 변수: composite checkpoint·online feature. 작동 원리: 각 seed의 exact logit을 CPU에서 수집한다."""

    outputs: list[torch.Tensor] = []
    exact_labels: tuple[str, ...] | None = None
    family_labels: tuple[str, ...] | None = None
    for base_path, adapter_path in zip(base_paths, adapter_paths, strict=True):
        model, adapter, base, _adapter_payload = _load_encoder06(
            base_path, adapter_path, device,
        )
        current_exact = tuple(str(value) for value in base["exact_labels"])
        current_family = tuple(str(value) for value in base["family_labels"])
        if exact_labels is not None and current_exact != exact_labels:
            raise ValueError("seed별 exact label 순서가 다릅니다.")
        exact_labels, family_labels = current_exact, current_family
        rows = []
        model.eval()
        adapter.eval()
        with torch.inference_mode():
            for start in range(0, len(features), batch_size):
                batch = features[start:start + batch_size].to(device)
                exact, _family = model.forward_online(adapter(batch))
                rows.append(exact.cpu())
        outputs.append(torch.cat(rows))
        del model, adapter
        if device.type == "cuda":
            torch.cuda.empty_cache()
    assert exact_labels is not None and family_labels is not None
    return outputs, exact_labels, family_labels


def main() -> None:
    """필요 변수: CLI 설정. 작동 원리: valid composite 세 seed의 합의 오류 보고서를 UTF-8 JSON으로 저장한다."""

    args = _parse_args()
    device = _resolve_device06(args.device)
    cache = torch.load(args.feature_cache, map_location="cpu", weights_only=True, mmap=True)
    features = cache["features"][:, 0].clone()
    targets = cache["targets"].long().clone()
    writers = cache.get("writers")
    if writers is None:
        writers = torch.full((len(targets),), -1, dtype=torch.long)
    else:
        writers = writers.long().clone()
    sources = _source_vector06(len(targets), args.source_counts, args.source_split)
    logits, labels, family_labels = _infer_logits06(
        args.base_checkpoint, args.adapter_checkpoint, features,
        device=device, batch_size=args.batch_size,
    )
    report = {
        "experiment": "MATH-INK-06-ONLINE-ERROR-CONSENSUS-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "device": str(device),
        "feature_cache": str(args.feature_cache),
        "base_checkpoints": [str(path) for path in args.base_checkpoint],
        "adapter_checkpoints": [str(path) for path in args.adapter_checkpoint],
        "label_count": len(labels),
        "family_count": len(family_labels),
        "metrics": summarize_online_consensus06(
            logits, targets, writers, labels, sources,
        ),
        "interpretation": {
            "all_seed_wrong": "세 seed가 모두 틀려 seed 증대만으로 회복되지 않는 데이터·표현 병목",
            "seed_oracle": "표본마다 정답을 낸 seed를 사후 선택한 비배포 상한",
            "product_validation": False,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()