"""세 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()