"""CROHME 연속 수식에서 AIFlow 0.6 행동 의미 역할 head를 GPU 학습한다.""" from __future__ import annotations import argparse from collections import Counter from copy import deepcopy from datetime import datetime, timezone import json import math from pathlib import Path import random import sys from typing import Any, Sequence import numpy as np import torch from torch import Tensor from torch.utils.data import DataLoader, TensorDataset 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.behavior_role_head06 import ( BEHAVIOR_CONTEXT_FEATURES06, BEHAVIOR_ROLE_LABELS06, BehaviorRoleHead06, behavior_role_index06, validate_behavior_context06, ) from math_grid_drawer.research.cross_visual import CROSS_FEATURE_NAMES, cross_pair_feature_rows from math_grid_drawer.research.ink06_canonical import canonicalize_ink06 from scripts.crohme_lattice_common import writer_fit_validation from scripts.train_crohme_segmentation_lattice_selector import _samples from scripts.train_math_ink_06_skeleton_adapter import ( _fused_exact_family_logits06, _paired_record_feature06, ) from scripts.audit_math_ink_06_case_context import _load_model06 from scripts.materialize_math_ink_06_paired_feature_cache import _selected_records TARGET_LABELS06 = frozenset({"c", "C", "x", "X", "z", "Z", r"\times"}) OPERATOR_LABELS06 = frozenset({ "+", "-", "=", "<", ">", "/", "*", r"\times", r"\div", r"\pm", r"\cdot", }) def _parse_args() -> argparse.Namespace: """필요 변수: adapter·공식 CROHME root·GPU 설정. 작동 원리: seed별 재현 가능한 행동학습 CLI를 만든다.""" parser = argparse.ArgumentParser(description="Train Math Ink 0.6 behavior role head") parser.add_argument("--adapter", type=Path, required=True) parser.add_argument( "--train-root", type=Path, default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/trainData", ) parser.add_argument( "--test-root", type=Path, default=PROJECT_ROOT / "research/data/R_noncommercial/ICFHR_package/CROHME2012_data/testDataGT", ) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--epochs", type=int, default=60) parser.add_argument("--patience", type=int, default=14) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--teacher-batch-size", type=int, default=256) parser.add_argument("--learning-rate", type=float, default=1e-3) parser.add_argument("--weight-decay", type=float, default=1e-3) parser.add_argument("--hidden", type=int, default=48) parser.add_argument("--dropout", type=float, default=0.15) parser.add_argument("--profile", default="median_height_32") parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda") parser.add_argument("--product-proxy-per-label", type=int, default=0) parser.add_argument("--product-proxy-loss-weight", type=float, default=0.35) parser.add_argument("--product-proxy-target-bases", default="cosuvwxz") parser.add_argument("--product-proxy-lowercase-ratio", type=float, default=1.0) parser.add_argument( "--data", type=Path, default=PROJECT_ROOT / "research/data/open_pretrain/hwrt_expanded_v2/hwrt_expanded.jsonl.gz", ) parser.add_argument( "--commercial-paired", type=Path, default=PROJECT_ROOT / "research/data/external_trajectory_v1/commercial_ccby4.jsonl.gz", ) parser.add_argument( "--dataset-registry", type=Path, default=PROJECT_ROOT / "research/dataset_registry.json", ) parser.add_argument( "--source-registry", type=Path, default=PROJECT_ROOT / "research/math_ink_06_source_registry.json", ) parser.add_argument( "--hwrt-approval", type=Path, default=PROJECT_ROOT / "research/approvals/HWRT-ODBL-USE-APPROVAL-v1.json", ) return parser.parse_args() def _seed06(seed: int) -> None: """필요 변수: 실험 seed. 작동 원리: Python·NumPy·PyTorch 초기화를 함께 고정한다.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _group_box06(group: frozenset[int], strokes: Sequence[dict[str, Any]]) -> dict[str, float]: """필요 변수: 정답 stroke group·원본 획. 작동 원리: 수식 좌표계의 bbox와 중심을 계산한다.""" points = [point for index in group for point in strokes[index]["points"]] xs, ys = [float(point[0]) for point in points], [float(point[1]) for point in points] left, top, right, bottom = min(xs), min(ys), max(xs), max(ys) return { "left": left, "top": top, "right": right, "bottom": bottom, "width": max(right - left, 1e-5), "height": max(bottom - top, 1e-5), "cx": (left + right) * 0.5, "cy": (top + bottom) * 0.5, } def _formula_box06(strokes: Sequence[dict[str, Any]]) -> dict[str, float]: """필요 변수: 한 수식의 전체 획. 작동 원리: canvas 정규화에 필요한 원점·크기를 반환한다.""" group = frozenset(range(len(strokes))) return _group_box06(group, strokes) def _canonical_group06( group: frozenset[int], strokes: Sequence[dict[str, Any]], formula_box: dict[str, float], ) -> Tensor: """필요 변수: symbol group·수식 bbox. 작동 원리: 원본 필순을 보존하며 formula-relative 128×19 입력을 만든다.""" selected = [] for order, stroke_index in enumerate(sorted(group)): stroke = strokes[stroke_index] selected.append({ "order": order, "stroke_id": order, "points": [ [ float(point[0]) - formula_box["left"], float(point[1]) - formula_box["top"], point[2] if len(point) > 2 else None, ] for point in stroke["points"] ], }) ink = canonicalize_ink06( selected, canvas_width=formula_box["width"], canvas_height=formula_box["height"], source_modality="online", trust_timestamps=False, ) return torch.from_numpy(ink.features) def _teacher_role_indices06(labels: Sequence[str]) -> dict[str, list[int]]: """필요 변수: 378 vocabulary. 작동 원리: 이웃의 예측 확률을 digit/identifier/operator 역할로 합칠 index를 만든다.""" digit = [index for index, label in enumerate(labels) if str(label).isdigit()] operator = [index for index, label in enumerate(labels) if str(label) in OPERATOR_LABELS06] identifier = [ index for index, label in enumerate(labels) if ( (len(str(label)) == 1 and str(label).isalpha()) or (str(label).startswith("\\") and str(label) not in OPERATOR_LABELS06) ) ] return {"digit": digit, "identifier": identifier, "operator": operator} def _role_mass06(probability: Tensor, indices: dict[str, list[int]]) -> list[float]: """필요 변수: 한 기호 teacher 확률·역할 index. 작동 원리: 네 coarse role 확률을 반환한다.""" digit = float(probability[indices["digit"]].sum()) identifier = float(probability[indices["identifier"]].sum()) operator = float(probability[indices["operator"]].sum()) return [digit, identifier, operator, max(0.0, 1.0 - digit - identifier - operator)] def _neighbor06( boxes: Sequence[dict[str, float]], index: int, *, direction: str, reference_height: float, ) -> int | None: """필요 변수: group bbox·기준 위치·방향. 작동 원리: 같은 로컬 행에 가까운 좌우 기호를 선택한다.""" target = boxes[index] candidates = [] for other_index, other in enumerate(boxes): if other_index == index: continue horizontal = other["cx"] < target["cx"] if direction == "left" else other["cx"] > target["cx"] if not horizontal: continue vertical = abs(other["cy"] - target["cy"]) / max(reference_height, target["height"], 1e-5) if vertical > 0.85: continue gap = ( max(0.0, target["left"] - other["right"]) if direction == "left" else max(0.0, other["left"] - target["right"]) ) candidates.append((gap + vertical * reference_height * 0.35, other_index)) return min(candidates)[1] if candidates else None def _family_log_probability06( log_probability: Tensor, labels: Sequence[str], base: str, ) -> tuple[float, float, float, float]: """필요 변수: teacher log 확률·case base. 작동 원리: lower/upper/times와 family 총질량을 공통 feature로 만든다.""" label_to_index = {str(label): index for index, label in enumerate(labels)} members = (base, base.upper(), r"\times") values = [ float(log_probability[label_to_index[label]]) if label in label_to_index else -20.0 for label in members ] available = [label_to_index[label] for label in members if label in label_to_index] mass = float(log_probability[available].logsumexp(dim=0).exp()) if available else 0.0 return values[0], values[1], values[2], mass def _context_row06( formula: dict[str, Any], index: int, boxes: Sequence[dict[str, float]], teacher_logits: Tensor, labels: Sequence[str], role_indices: dict[str, list[int]], ) -> list[float]: """필요 변수: 수식·target group·teacher logits. 작동 원리: 누수 없는 예측 이웃과 실제 배치·획 feature를 49차로 만든다.""" target_label = str(formula["truth_labels"][index]) base = "x" if target_label == r"\times" else target_label.lower() probability = teacher_logits.softmax(dim=1) log_probability = teacher_logits.log_softmax(dim=1) target_probability = probability[index] lower, upper, times, family_mass = _family_log_probability06( log_probability[index], labels, base, ) entropy = float( -(target_probability * target_probability.clamp_min(1e-9).log()).sum() / max(math.log(len(labels)), 1.0) ) heights = [box["height"] for box in boxes] reference_height = float(np.median(heights)) widths = [box["width"] for box in boxes] reference_width = float(np.median(widths)) left_index = _neighbor06(boxes, index, direction="left", reference_height=reference_height) right_index = _neighbor06(boxes, index, direction="right", reference_height=reference_height) left_role = _role_mass06(probability[left_index], role_indices) if left_index is not None else [0.0] * 4 right_role = _role_mass06(probability[right_index], role_indices) if right_index is not None else [0.0] * 4 target = boxes[index] left = boxes[left_index] if left_index is not None else None right = boxes[right_index] if right_index is not None else None formula_box = _formula_box06(formula["strokes"]) group = formula["truth_groups"][index] target_strokes = [formula["strokes"][stroke_index] for stroke_index in sorted(group)] cross = np.zeros(len(CROSS_FEATURE_NAMES), dtype=np.float32) if len(target_strokes) == 2: cross = cross_pair_feature_rows([(0, 1)], target_strokes)[0] point_count = sum(len(stroke.get("points") or []) for stroke in target_strokes) values = [ lower, upper, times, family_mass, float(target_probability.max()), entropy, *left_role, *right_role, max(0.0, target["left"] - left["right"]) / reference_height if left else 2.0, max(0.0, right["left"] - target["right"]) / reference_height if right else 2.0, (target["cy"] - left["cy"]) / reference_height if left else 0.0, (right["cy"] - target["cy"]) / reference_height if right else 0.0, (target["cx"] - formula_box["left"]) / formula_box["width"], (target["cy"] - formula_box["top"]) / formula_box["height"], target["width"] / formula_box["width"], target["height"] / formula_box["height"], target["height"] / max(reference_height, 1e-5), target["width"] / max(reference_width, 1e-5), min(len(group) / 8.0, 1.0), min(point_count / 256.0, 1.0), float(base == "c"), float(base == "x"), float(base == "z"), *cross.tolist(), float(len(target_strokes) == 2), ] validate_behavior_context06(values) return values def _materialize_split06( samples: Sequence[dict[str, Any]], engine, adapter: torch.nn.Module, *, device: torch.device, teacher_batch_size: int, return_metadata: bool = False, ) -> tuple[TensorDataset, dict[str, int]] | tuple[TensorDataset, dict[str, int], list[dict[str, Any]]]: """필요 변수: 수식 split·frozen teacher. 작동 원리: 모든 이웃은 teacher 예측으로 만들고 target stroke만 corpus에 보존한다.""" labels = tuple(str(label) for label in engine.labels) role_indices = _teacher_role_indices06(labels) sequence_rows: list[Tensor] = [] context_rows: list[list[float]] = [] target_rows: list[int] = [] metadata_rows: list[dict[str, Any]] = [] counts: Counter[str] = Counter() engine.model.eval() adapter.eval() with torch.inference_mode(): for formula in samples: groups = formula["truth_groups"] boxes = [_group_box06(group, formula["strokes"]) for group in groups] formula_box = _formula_box06(formula["strokes"]) sequences = [ _canonical_group06(group, formula["strokes"], formula_box) for group in groups ] logits_parts = [] for start in range(0, len(sequences), teacher_batch_size): batch = torch.stack(sequences[start:start + teacher_batch_size]).to(device) four_hypotheses = batch.unsqueeze(1).expand(-1, 4, -1, -1).contiguous() logits, _family = _fused_exact_family_logits06(engine, adapter, four_hypotheses) logits_parts.append(logits.cpu()) teacher_logits = torch.cat(logits_parts) for index, label_value in enumerate(formula["truth_labels"]): label = str(label_value) if label not in TARGET_LABELS06: continue target = behavior_role_index06(label) if target is None: continue sequence_rows.append(sequences[index]) context_rows.append(_context_row06( formula, index, boxes, teacher_logits, labels, role_indices, )) target_rows.append(target) counts[BEHAVIOR_ROLE_LABELS06[target]] += 1 if return_metadata: teacher_index = int(teacher_logits[index].argmax()) metadata_rows.append({ "sample_id": str(formula["sample_id"]), "truth_label": label, "teacher_label": labels[teacher_index], "teacher_probability": float( teacher_logits[index].softmax(dim=0)[teacher_index] ), }) if not sequence_rows: raise ValueError("행동 role 학습 표본이 없습니다.") dataset = TensorDataset( torch.stack(sequence_rows), torch.tensor(context_rows, dtype=torch.float32), torch.tensor(target_rows, dtype=torch.long), ) if return_metadata: return dataset, dict(counts), metadata_rows return dataset, dict(counts) def _synthetic_scale06(label: str, rng: np.random.Generator) -> float: """필요 변수: case label·seed RNG. 작동 원리: 실제 행의 겹치는 크기 분포를 포함한 수식 배치 높이를 만든다.""" ambiguous = bool(rng.random() < 0.20) if label.isupper(): value = rng.normal(0.84 if ambiguous else 0.96, 0.055) else: value = rng.normal(0.84 if ambiguous else 0.68, 0.070) return float(np.clip(value, 0.48, 1.05)) def _place_synthetic_row06(feature: Tensor, scale: float) -> Tensor: """필요 변수: 실제 glyph 128×19 feature·합성 높이. 작동 원리: 필순은 보존하고 row-relative canvas/baseline만 배치한다.""" output = feature.clone() top = (1.0 - scale) * 0.5 bottom = top + scale output[:, 3] = top + output[:, 1] * scale output[:, 10] = top output[:, 11] = bottom output[:, 12] = scale output[:, 13] = 0.5 output[:, 14] = 1.0 return output def _product_proxy_records06(args: argparse.Namespace, labels: Sequence[str]) -> list[dict[str, Any]]: """필요 변수: fail-closed registry·teacher vocabulary. 작동 원리: 승인 training split만 읽어 대소문자 proxy 원천을 만든다.""" selection_args = argparse.Namespace( data=args.data, commercial_paired=args.commercial_paired, dataset_registry=args.dataset_registry, source_registry=args.source_registry, hwrt_approval=args.hwrt_approval, paired_source=["uci-uji-pen-v1", "uci-uji-pen-v2"], split="training", maximum_paired_train_per_label=args.product_proxy_per_label, maximum_paired_eval_per_label=0, ) records, _fingerprint = _selected_records(selection_args, list(labels)) return records def _materialize_product_proxy06( records: Sequence[dict[str, Any]], engine, adapter: torch.nn.Module, *, seed: int, device: torch.device, teacher_batch_size: int, target_bases: frozenset[str], lowercase_ratio: float, ) -> tuple[TensorDataset, dict[str, Any]]: """필요 변수: 승인 고립 trajectory·frozen teacher. 작동 원리: 실제 stroke를 확률적 수식 행에 배치해 case 행동 proxy를 만든다.""" labels = tuple(str(label) for label in engine.labels) role_indices = _teacher_role_indices06(labels) rng = np.random.default_rng(seed + 600) if not target_bases or not target_bases <= frozenset({"c", "o", "s", "u", "v", "w", "x", "z"}): raise ValueError("제품 proxy target base가 지원 case family와 다릅니다.") if not 0.0 <= lowercase_ratio <= 1.0: raise ValueError("제품 proxy lowercase ratio는 0~1이어야 합니다.") targets = [ record for record in records if ( len(str(record["label"])) == 1 and str(record["label"]).lower() in target_bases ) ] targets = [ record for index, record in enumerate(targets) if ( str(record["label"]).isupper() or (index * 2654435761 + seed) % 10_000 < int(lowercase_ratio * 10_000) ) ] anchors = [record for record in records if str(record["label"]).isdigit()] if not targets or not anchors: raise ValueError("제품 proxy에 case target 또는 숫자 anchor가 없습니다.") anchors_by_writer: dict[tuple[str, str], list[dict[str, Any]]] = {} anchors_by_source: dict[str, list[dict[str, Any]]] = {} for record in anchors: key = (str(record.get("source")), str(record.get("writer_key"))) anchors_by_writer.setdefault(key, []).append(record) anchors_by_source.setdefault(str(record.get("source")), []).append(record) sequences: list[Tensor] = [] triplet_features: list[Tensor] = [] metadata: list[dict[str, Any]] = [] for index, record in enumerate(targets): key = (str(record.get("source")), str(record.get("writer_key"))) pool = anchors_by_writer.get(key) or anchors_by_source.get(str(record.get("source"))) or anchors left_record = pool[(index * 2) % len(pool)] right_record = pool[(index * 2 + 1) % len(pool)] label = str(record["label"]) scale = _synthetic_scale06(label, rng) target_feature = _place_synthetic_row06(_paired_record_feature06(record)[0], scale) left_scale = float(np.clip(rng.normal(0.96, 0.04), 0.82, 1.05)) right_scale = float(np.clip(rng.normal(0.96, 0.04), 0.82, 1.05)) left_feature = _place_synthetic_row06(_paired_record_feature06(left_record)[0], left_scale) right_feature = _place_synthetic_row06(_paired_record_feature06(right_record)[0], right_scale) sequences.append(target_feature) triplet_features.extend((target_feature, left_feature, right_feature)) metadata.append({ "record": record, "label": label, "scale": scale, "left_scale": left_scale, "right_scale": right_scale, "gap_left": float(rng.uniform(0.12, 0.55)), "gap_right": float(rng.uniform(0.12, 0.55)), }) teacher_rows = [] engine.model.eval() adapter.eval() with torch.inference_mode(): for start in range(0, len(triplet_features), teacher_batch_size): batch = torch.stack(triplet_features[start:start + teacher_batch_size]).to(device) hypotheses = batch.unsqueeze(1).expand(-1, 4, -1, -1).contiguous() logits, _family = _fused_exact_family_logits06(engine, adapter, hypotheses) teacher_rows.append(logits.cpu()) teacher_logits = torch.cat(teacher_rows).reshape(len(targets), 3, -1) context_rows: list[list[float]] = [] target_rows: list[int] = [] counts: Counter[str] = Counter() ambiguous_rows = 0 for index, row in enumerate(metadata): label = row["label"] target = behavior_role_index06(label) if target is None: continue probability = teacher_logits[index].softmax(dim=1) log_probability = teacher_logits[index, 0].log_softmax(dim=0) base = label.lower() lower, upper, times, family_mass = _family_log_probability06( log_probability, labels, base, ) target_probability = probability[0] entropy = float( -(target_probability * target_probability.clamp_min(1e-9).log()).sum() / max(math.log(len(labels)), 1.0) ) left_role = _role_mass06(probability[1], role_indices) right_role = _role_mass06(probability[2], role_indices) record = row["record"] target_strokes = list(record["strokes"]) cross = np.zeros(len(CROSS_FEATURE_NAMES), dtype=np.float32) if len(target_strokes) == 2: cross = cross_pair_feature_rows([(0, 1)], target_strokes)[0] point_count = sum(len(stroke.get("points") or []) for stroke in target_strokes) scale = float(row["scale"]) reference_height = (float(row["left_scale"]) + float(row["right_scale"])) * 0.5 aspect = float(sequences[index][0, 9]) values = [ lower, upper, times, family_mass, float(target_probability.max()), entropy, *left_role, *right_role, float(row["gap_left"]), float(row["gap_right"]), 0.0, 0.0, 0.5, 0.5, min(0.24 * max(aspect, 0.2), 0.45), scale, scale / max(reference_height, 1e-5), max(aspect, 1e-3), min(len(target_strokes) / 8.0, 1.0), min(point_count / 256.0, 1.0), float(base == "c"), float(base == "x"), float(base == "z"), *cross.tolist(), float(len(target_strokes) == 2), ] validate_behavior_context06(values) context_rows.append(values) target_rows.append(target) counts[BEHAVIOR_ROLE_LABELS06[target]] += 1 ambiguous_rows += int( (label.isupper() and scale < 0.90) or (label.islower() and scale > 0.78) ) return TensorDataset( torch.stack(sequences), torch.tensor(context_rows, dtype=torch.float32), torch.tensor(target_rows, dtype=torch.long), ), { "samples": len(target_rows), "role_counts": dict(counts), "ambiguous_size_samples": ambiguous_rows, "sources": dict(Counter(str(record.get("source")) for record in targets)), "track": "P_approved_synthetic_layout_proxy", "actual_continuous_formula": False, "target_bases": sorted(target_bases), "lowercase_ratio": lowercase_ratio, } def _with_weight06(dataset: TensorDataset, weight: float) -> TensorDataset: """필요 변수: dataset·loss weight. 작동 원리: source별 기여도를 각 행에 명시적으로 붙인다.""" return TensorDataset( *dataset.tensors, torch.full((len(dataset),), float(weight), dtype=torch.float32), ) def _combine_training06(parts: Sequence[TensorDataset]) -> TensorDataset: """필요 변수: 동일 계약의 weighted dataset. 작동 원리: source 순서를 보존해 한 학습 tensor로 결합한다.""" return TensorDataset(*( torch.cat([dataset.tensors[index] for dataset in parts], dim=0) for index in range(len(parts[0].tensors)) )) def _normalize_context06( training: TensorDataset, validation: TensorDataset, testing: TensorDataset, *, statistics_reference: TensorDataset | None = None, ) -> tuple[TensorDataset, TensorDataset, TensorDataset, Tensor, Tensor]: """필요 변수: 세 split context. 작동 원리: fit 통계만 사용해 validation/test 누수를 차단한다.""" reference = statistics_reference or training mean = reference.tensors[1].mean(dim=0) scale = reference.tensors[1].std(dim=0).clamp_min(1e-5) def normalized(dataset: TensorDataset) -> TensorDataset: """필요 변수: 원본 dataset. 작동 원리: sequence/target은 보존하고 context만 표준화한다.""" return TensorDataset( dataset.tensors[0], (dataset.tensors[1] - mean) / scale, *dataset.tensors[2:], ) return normalized(training), normalized(validation), normalized(testing), mean, scale def _metrics06(logits: Tensor, target: Tensor) -> dict[str, Any]: """필요 변수: N×3 logits·정답. 작동 원리: accuracy·macro-F1·class recall·ECE를 같은 분모로 계산한다.""" prediction = logits.argmax(dim=1) probability = logits.softmax(dim=1) confusion = torch.zeros((3, 3), dtype=torch.long) for truth, selected in zip(target.tolist(), prediction.tolist(), strict=True): confusion[truth, selected] += 1 f1_values, recalls = [], {} for index, label in enumerate(BEHAVIOR_ROLE_LABELS06): true_positive = int(confusion[index, index]) false_positive = int(confusion[:, index].sum()) - true_positive false_negative = int(confusion[index, :].sum()) - true_positive precision = true_positive / max(true_positive + false_positive, 1) recall = true_positive / max(true_positive + false_negative, 1) f1_values.append(2 * precision * recall / max(precision + recall, 1e-12)) recalls[label] = recall confidence = probability.max(dim=1).values correct = prediction.eq(target) ece = 0.0 for lower in torch.arange(0.0, 1.0, 0.1): selected = (confidence >= lower) & (confidence < lower + 0.1) if selected.any(): ece += float(selected.float().mean()) * abs( float(correct[selected].float().mean()) - float(confidence[selected].mean()) ) return { "samples": len(target), "accuracy": float(correct.float().mean()), "macro_f1": float(np.mean(f1_values)), "recall": recalls, "confusion": confusion.tolist(), "ece": ece, } def _evaluate06( model: BehaviorRoleHead06, dataset: TensorDataset, *, batch_size: int, device: torch.device, ) -> tuple[dict[str, Any], Tensor, Tensor]: """필요 변수: behavior head·split. 작동 원리: 순서를 보존해 전체 logits와 정답을 모은다.""" model.eval() logits_rows, target_rows = [], [] with torch.inference_mode(): for batch in DataLoader(dataset, batch_size=batch_size, shuffle=False): sequence, context, target = batch[:3] logits_rows.append(model(sequence.to(device), context.to(device)).cpu()) target_rows.append(target) logits, targets = torch.cat(logits_rows), torch.cat(target_rows) return _metrics06(logits, targets), logits, targets def _teacher_baseline06(dataset: TensorDataset) -> dict[str, Any]: """필요 변수: 정규화 전 context의 teacher logp 3개. 작동 원리: 행동 head가 실제로 teacher family를 개선했는지 비교한다.""" logits = dataset.tensors[1][:, :3] return _metrics06(logits, dataset.tensors[2]) def main() -> None: """필요 변수: seed별 adapter와 CROHME 공식 split. 작동 원리: validation 선택 후 official test를 한 번 평가한다.""" args = _parse_args() if args.train_root.name.casefold() != "traindata" or args.test_root.name.casefold() != "testdatagt": raise ValueError("CROHME2012 공식 trainData/testDataGT 조합만 허용합니다.") device = torch.device(args.device) if device.type == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA 학습을 요청했지만 사용할 수 없습니다.") _seed06(args.seed) adapter_payload = torch.load(args.adapter, map_location="cpu", weights_only=False) base_checkpoint = Path(str(adapter_payload["base_checkpoint"])) if not base_checkpoint.is_absolute(): base_checkpoint = PROJECT_ROOT / base_checkpoint engine, adapter = _load_model06(base_checkpoint, args.adapter, device) fit_samples, validation_samples = writer_fit_validation(args.train_root, args.profile) test_samples = _samples(args.test_root, args.profile) training_raw, training_counts = _materialize_split06( fit_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size, ) validation_raw, validation_counts = _materialize_split06( validation_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size, ) testing_raw, testing_counts = _materialize_split06( test_samples, engine, adapter, device=device, teacher_batch_size=args.teacher_batch_size, ) proxy_info: dict[str, Any] | None = None training_weighted = _with_weight06(training_raw, 1.0) if args.product_proxy_per_label > 0: product_records = _product_proxy_records06(args, engine.labels) proxy_raw, proxy_info = _materialize_product_proxy06( product_records, engine, adapter, seed=args.seed, device=device, teacher_batch_size=args.teacher_batch_size, target_bases=frozenset(args.product_proxy_target_bases), lowercase_ratio=args.product_proxy_lowercase_ratio, ) training_weighted = _combine_training06(( training_weighted, _with_weight06(proxy_raw, args.product_proxy_loss_weight), )) del engine, adapter if device.type == "cuda": torch.cuda.empty_cache() teacher_baseline = { "validation": _teacher_baseline06(validation_raw), "official_test": _teacher_baseline06(testing_raw), } training, validation, testing, mean, scale = _normalize_context06( training_weighted, validation_raw, testing_raw, statistics_reference=training_raw, ) model = BehaviorRoleHead06(hidden=args.hidden, dropout=args.dropout).to(device) counts = torch.bincount(training.tensors[2], minlength=3).float() class_weight = (counts.sum() / counts.clamp_min(1.0)).sqrt() class_weight = (class_weight / class_weight.mean()).to(device) optimizer = torch.optim.AdamW( model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay, ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) loader = DataLoader( training, batch_size=args.batch_size, shuffle=True, generator=torch.Generator().manual_seed(args.seed), ) best_key = (-1.0, -1.0) best_state: dict[str, Tensor] | None = None best_epoch = 0 history = [] stale = 0 for epoch in range(1, args.epochs + 1): model.train() total_loss = 0.0 samples = 0 for sequence, context, target, sample_weight in loader: sequence, context = sequence.to(device), context.to(device) target, sample_weight = target.to(device), sample_weight.to(device) optimizer.zero_grad(set_to_none=True) logits = model(sequence, context) losses = torch.nn.functional.cross_entropy( logits, target, weight=class_weight, reduction="none", ) loss = (losses * sample_weight).sum() / sample_weight.sum().clamp_min(1e-8) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0) optimizer.step() total_loss += float(loss.detach()) * len(target) samples += len(target) scheduler.step() validation_metrics, _logits, _targets = _evaluate06( model, validation, batch_size=args.batch_size, device=device, ) row = { "epoch": epoch, "training_loss": total_loss / max(samples, 1), "learning_rate": optimizer.param_groups[0]["lr"], "validation": validation_metrics, } history.append(row) key = (validation_metrics["macro_f1"], validation_metrics["accuracy"]) if key > best_key: best_key = key best_epoch = epoch best_state = deepcopy({key: value.detach().cpu() for key, value in model.state_dict().items()}) stale = 0 else: stale += 1 if epoch == 1 or epoch % 5 == 0: print(json.dumps({ "seed": args.seed, "epoch": epoch, "loss": row["training_loss"], "validation_accuracy": validation_metrics["accuracy"], "validation_macro_f1": validation_metrics["macro_f1"], }, ensure_ascii=False), flush=True) if stale >= args.patience: break if best_state is None: raise RuntimeError("behavior role checkpoint가 선택되지 않았습니다.") model.load_state_dict(best_state) validation_metrics, _validation_logits, _validation_targets = _evaluate06( model, validation, batch_size=args.batch_size, device=device, ) test_metrics, _test_logits, _test_targets = _evaluate06( model, testing, batch_size=args.batch_size, device=device, ) args.output.mkdir(parents=True, exist_ok=True) checkpoint = args.output / "behavior_role_head.pt" torch.save({ "schema": "aiflow-math-ink-06-behavior-role-v1", "state_dict": best_state, "context_mean": mean, "context_scale": scale, "context_features": BEHAVIOR_CONTEXT_FEATURES06, "role_labels": BEHAVIOR_ROLE_LABELS06, "sequence_channels": 19, "hidden": args.hidden, "dropout": args.dropout, "seed": args.seed, "teacher_adapter": str(args.adapter), "selected_epoch": best_epoch, "track": "R_noncommercial_only", "product_validation": False, }, checkpoint) report = { "experiment": "R-MATH-INK-06-BEHAVIOR-ROLE-001", "generated_at": datetime.now(timezone.utc).isoformat(), "seed": args.seed, "device": str(device), "cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None, "teacher_adapter": str(args.adapter), "base_checkpoint": str(base_checkpoint), "split_contract": "CROHME2012 trainData writer fit/validation; testDataGT official held-out", "formulas": { "fit": len(fit_samples), "validation": len(validation_samples), "official_test": len(test_samples), }, "role_samples": { "fit": training_counts, "validation": validation_counts, "official_test": testing_counts, }, "product_proxy": proxy_info, "product_proxy_loss_weight": ( args.product_proxy_loss_weight if proxy_info is not None else 0.0 ), "teacher_baseline": teacher_baseline, "selected_epoch": best_epoch, "validation": validation_metrics, "official_test": test_metrics, "history": history, "checkpoint": checkpoint.name, "checkpoint_bytes": checkpoint.stat().st_size, "track": "R_noncommercial_only", "product_validation": False, "interpretation_limit": ( "CROHME 연구용 연속 수식 행동 head이며 상용 checkpoint에 병합할 수 없다. " "P-track writer/device-disjoint 재학습이 필요하다." ), } (args.output / "report.json").write_text( json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(json.dumps({ "seed": args.seed, "selected_epoch": best_epoch, "teacher_test": teacher_baseline["official_test"], "behavior_test": test_metrics, "checkpoint": str(checkpoint), "product_validation": False, }, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()