scripts/calibrate_math_ink_06_online_family_fusion.py

"""Validation에서 exact/family logit 결합 가중치를 고정하고 paired-test에 한 번 적용한다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
import math
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 scripts.train_math_ink_06_p_boundary_auxiliary import _load_encoder06
from scripts.train_math_ink_06_skeleton_adapter import _resolve_device06
from math_grid_drawer.research.trajectory_sequence import shape_family


def _parse_args() -> argparse.Namespace:
    """필요 변수: validation/test cache와 세 seed. 작동 원리: test 선택을 금지한 calibration CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Calibrate Math Ink 0.6 online family fusion")
    parser.add_argument("--validation-cache", type=Path, required=True)
    parser.add_argument("--test-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("--weights", default="0,0.025,0.05,0.075,0.1,0.125,0.15,0.2,0.25,0.3")
    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("fusion calibration에는 seed 두 개 이상이 필요합니다.")
    return args


def _macro_f106(targets: torch.Tensor, predictions: torch.Tensor) -> float:
    """필요 변수: 정답·예측 index. 작동 원리: test에 없는 class를 분모에서 제외한 macro-F1을 계산한다."""

    values = []
    for label in targets.unique().tolist():
        truth = targets.eq(label)
        predicted = predictions.eq(label)
        true_positive = int((truth & predicted).sum())
        denominator = 2 * true_positive + int((truth & ~predicted).sum()) + int((~truth & predicted).sum())
        values.append(2 * true_positive / denominator if denominator else 0.0)
    return sum(values) / max(len(values), 1)


def _metrics06(logits: torch.Tensor, targets: torch.Tensor) -> dict[str, float | int]:
    """필요 변수: fused logit·정답. 작동 원리: 동일 분모의 top-1/top-5/macro-F1을 반환한다."""

    prediction = logits.argmax(dim=-1)
    top5 = logits.topk(min(5, logits.shape[-1]), dim=-1).indices
    return {
        "samples": len(targets),
        "top1": float(prediction.eq(targets).float().mean()),
        "top5": float(top5.eq(targets[:, None]).any(dim=-1).float().mean()),
        "macro_f1": _macro_f106(targets, prediction),
    }


def fusion_sweep06(
    exact_by_seed: Sequence[torch.Tensor],
    family_by_seed: Sequence[torch.Tensor],
    targets: torch.Tensor,
    exact_family_index: torch.Tensor,
    weights: Sequence[float],
) -> list[dict[str, float | int]]:
    """필요 변수: seed별 exact/family logit·가중치. 작동 원리: 확률공간 seed ensemble을 weight별 평가한다."""

    if len(exact_by_seed) != len(family_by_seed) or not exact_by_seed:
        raise ValueError("exact/family seed 출력 개수가 올바르지 않습니다.")
    rows = []
    for weight in weights:
        seed_joint = []
        for exact, family in zip(exact_by_seed, family_by_seed, strict=True):
            joint = exact.log_softmax(dim=-1)
            if weight:
                joint = joint + float(weight) * family.log_softmax(dim=-1)[:, exact_family_index]
            seed_joint.append(joint)
        ensemble = torch.logsumexp(torch.stack(seed_joint), dim=0) - math.log(len(seed_joint))
        rows.append({"family_fusion_weight": float(weight), **_metrics06(ensemble, targets)})
    return rows


def _infer_split06(
    cache_path: Path,
    base_paths: Sequence[Path],
    adapter_paths: Sequence[Path],
    *,
    device: torch.device,
    batch_size: int,
) -> tuple[list[torch.Tensor], list[torch.Tensor], torch.Tensor, torch.Tensor]:
    """필요 변수: split cache·composite seed. 작동 원리: fusion 전 exact/family logit과 사상을 수집한다."""

    cache = torch.load(cache_path, map_location="cpu", weights_only=True, mmap=True)
    features = cache["features"][:, 0]
    targets = cache["targets"].long().clone()
    exact_rows, family_rows = [], []
    family_index: torch.Tensor | 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)
        family_to_index = {
            str(label): index for index, label in enumerate(base["family_labels"])
        }
        current_family_index = torch.tensor([
            family_to_index[shape_family(str(label))]
            for label in base["exact_labels"]
        ], dtype=torch.long)
        if family_index is not None and not torch.equal(family_index, current_family_index):
            raise ValueError("seed별 exact→family 사상이 다릅니다.")
        family_index = current_family_index
        exact_batches, family_batches = [], []
        model.eval()
        adapter.eval()
        with torch.inference_mode():
            for start in range(0, len(features), batch_size):
                exact, family = model.forward_online(
                    adapter(features[start:start + batch_size].to(device)),
                )
                exact_batches.append(exact.cpu())
                family_batches.append(family.cpu())
        exact_rows.append(torch.cat(exact_batches))
        family_rows.append(torch.cat(family_batches))
        del model, adapter
        if device.type == "cuda":
            torch.cuda.empty_cache()
    assert family_index is not None
    return exact_rows, family_rows, targets, family_index


def main() -> None:
    """필요 변수: CLI 설정. 작동 원리: validation winner만 test에 적용하고 결과를 UTF-8 JSON으로 고정한다."""

    args = _parse_args()
    weights = tuple(float(value.strip()) for value in args.weights.split(",") if value.strip())
    if not weights or any(weight < 0.0 or weight > 1.0 for weight in weights):
        raise ValueError("family fusion weight는 0~1 범위여야 합니다.")
    device = _resolve_device06(args.device)
    validation = _infer_split06(
        args.validation_cache, args.base_checkpoint, args.adapter_checkpoint,
        device=device, batch_size=args.batch_size,
    )
    validation_sweep = fusion_sweep06(*validation, weights)
    selected = max(
        validation_sweep,
        key=lambda row: (float(row["top1"]), float(row["macro_f1"]), -float(row["family_fusion_weight"])),
    )
    test = _infer_split06(
        args.test_cache, args.base_checkpoint, args.adapter_checkpoint,
        device=device, batch_size=args.batch_size,
    )
    test_result = fusion_sweep06(
        *test, (float(selected["family_fusion_weight"]),),
    )[0]
    zero_test = fusion_sweep06(*test, (0.0,))[0]
    report = {
        "experiment": "MATH-INK-06-ONLINE-FAMILY-FUSION-CALIBRATION-001",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "selection_contract": "validation only; paired-test evaluated once after weight lock",
        "device": str(device),
        "validation_sweep": validation_sweep,
        "selected_validation": selected,
        "paired_test_zero_weight": zero_test,
        "paired_test_selected_weight": test_result,
        "paired_test_gain_pp": {
            "top1": (float(test_result["top1"]) - float(zero_test["top1"])) * 100.0,
            "top5": (float(test_result["top5"]) - float(zero_test["top5"])) * 100.0,
            "macro_f1": (float(test_result["macro_f1"]) - float(zero_test["macro_f1"])) * 100.0,
        },
        "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()