"""P-track exact/family rehearsal과 formula-boundary loss로 shared encoder 일부를 공동 파인튜닝한다.""" from __future__ import annotations import argparse from datetime import datetime, timezone import json from pathlib import Path import sys import numpy as np import torch 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.math_ink_06 import boundary_auxiliary_loss06 from math_grid_drawer.research.trajectory_sequence import shape_family from scripts.train_math_ink_06_p_boundary_auxiliary import ( _balanced_boundary_set06, _file_sha25606, _load_encoder06, _load_feature_cache06, _metrics06, _select_threshold06, ) from scripts.train_math_ink_06_skeleton_adapter import _resolve_device06 def _parse_args() -> argparse.Namespace: """필요 변수: P cache·seed checkpoint·초기 boundary head. 작동 원리: 제한된 shared joint fine-tuning CLI를 만든다.""" parser = argparse.ArgumentParser(description="Train Math Ink 0.6 P boundary joint model") parser.add_argument( "--training-cache", type=Path, default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-training-e811a2cfb9871e990f87.pt"), ) parser.add_argument( "--validation-cache", type=Path, default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-validation-b98e59caaacf15025b4f.pt"), ) parser.add_argument( "--test-cache", type=Path, default=Path(r"D:\Aiflow-CUDA\ink06_feature_cache\paired-paired-test-26bee2f320c6f0a7eca3.pt"), ) parser.add_argument( "--base-checkpoint", type=Path, default=PROJECT_ROOT / "research/runs/math_ink_06_federated_virtual_ce025_family010_seed17_20260723/math_ink_06_candidate.pt", ) parser.add_argument( "--adapter-checkpoint", type=Path, default=PROJECT_ROOT / "research/runs/math_ink_06_online_casecontext_refined_seed17_20260723/skeleton_adapter.pt", ) parser.add_argument( "--boundary-checkpoint", type=Path, default=PROJECT_ROOT / "research/runs/math_ink_06_p_boundary_auxiliary_layout5_20260724/boundary_auxiliary_head.pt", ) parser.add_argument("--samples-per-class", type=int, default=1200) parser.add_argument("--rehearsal-samples", type=int, default=4800) parser.add_argument("--epochs", type=int, default=8) parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--head-learning-rate", type=float, default=5e-4) parser.add_argument("--classifier-learning-rate", type=float, default=1e-5) parser.add_argument("--shared-learning-rate", type=float, default=5e-6) parser.add_argument("--boundary-loss-weight", type=float, default=0.5) parser.add_argument("--authentic-negative-weight", type=float, default=1.5) parser.add_argument("--maximum-exact-regression-pp", type=float, default=0.5) parser.add_argument("--maximum-family-regression-pp", type=float, default=0.5) parser.add_argument("--seed", type=int, default=17) parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto") parser.add_argument("--output", type=Path, required=True) return parser.parse_args() def _subset06( features: torch.Tensor, targets: torch.Tensor, *, count: int, seed: int, ) -> tuple[torch.Tensor, torch.Tensor]: """필요 변수: authentic cache·상한·seed. 작동 원리: 결정적 무복원 subset으로 exact/family rehearsal 분모를 만든다.""" if count <= 0: return features, targets generator = torch.Generator().manual_seed(seed) indices = torch.randperm(len(features), generator=generator)[:min(count, len(features))] return features[indices].clone(), targets[indices].clone() def _family_target_map06(base: dict) -> torch.Tensor: """필요 변수: exact/family vocabulary. 작동 원리: authentic exact index를 visual-family index로 변환한다.""" family_index = {str(label): index for index, label in enumerate(base["family_labels"])} return torch.tensor([ family_index[shape_family(str(label))] for label in base["exact_labels"] ], dtype=torch.long) def _authentic_metrics06( model, adapter, features: torch.Tensor, targets: torch.Tensor, family_map: torch.Tensor, *, device: torch.device, batch_size: int, ) -> dict[str, float]: """필요 변수: authentic online feature·exact target. 작동 원리: exact/family top-1을 고정 encoder 경로에서 평가한다.""" exact_match = family_match = total = 0 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) target = targets[start:start + batch_size] exact, family = model.forward_online(adapter(batch)) exact_match += int((exact.argmax(dim=-1).cpu() == target).sum()) family_target = family_map[target] family_match += int((family.argmax(dim=-1).cpu() == family_target).sum()) total += len(target) return { "samples": total, "exact_top1": exact_match / max(total, 1), "family_top1": family_match / max(total, 1), } def _boundary_logits06( model, adapter, features: torch.Tensor, *, device: torch.device, batch_size: int, ) -> torch.Tensor: """필요 변수: formula candidate feature. 작동 원리: 선택적 boundary head logit만 batch 추론해 CPU로 반환한다.""" rows = [] model.eval() adapter.eval() with torch.inference_mode(): for start in range(0, len(features), batch_size): batch = adapter(features[start:start + batch_size].to(device)) _exact, _family, boundary = model.forward_online_with_boundary(batch) rows.append(boundary.cpu()) return torch.cat(rows) def _selected_state06(model) -> dict[str, torch.Tensor]: """필요 변수: 공동학습 모델. 작동 원리: 실제로 열린 module만 portable delta checkpoint로 추출한다.""" prefixes = ( "trajectory_encoder.blocks.3.", "trajectory_encoder.attention.", "exact_head.", "family_head.", "boundary_head.", ) return { key: value.detach().cpu().clone() for key, value in model.state_dict().items() if key.startswith(prefixes) } def main() -> None: """필요 변수: P train/validation/test. 작동 원리: validation 회귀 gate 안 joint winner만 paired test에 한 번 적용한다.""" args = _parse_args() torch.manual_seed(args.seed) device = _resolve_device06(args.device) train_features, train_targets, train_cache_key = _load_feature_cache06(args.training_cache) validation_features, validation_targets, validation_cache_key = _load_feature_cache06(args.validation_cache) test_features, test_targets, test_cache_key = _load_feature_cache06(args.test_cache) train_rehearsal_x, train_rehearsal_y = _subset06( train_features, train_targets, count=args.rehearsal_samples, seed=args.seed, ) validation_rehearsal_x, validation_rehearsal_y = _subset06( validation_features, validation_targets, count=0, seed=args.seed + 1, ) test_rehearsal_x, test_rehearsal_y = _subset06( test_features, test_targets, count=0, seed=args.seed + 2, ) train_boundary_x, train_boundary_y = _balanced_boundary_set06( train_features, train_targets, samples_per_class=args.samples_per_class, seed=args.seed, ) validation_boundary_x, validation_boundary_y = _balanced_boundary_set06( validation_features, validation_targets, samples_per_class=args.samples_per_class, seed=args.seed + 1, ) test_boundary_x, test_boundary_y = _balanced_boundary_set06( test_features, test_targets, samples_per_class=args.samples_per_class, seed=args.seed + 2, ) model, adapter, base, _adapter_payload = _load_encoder06( args.base_checkpoint, args.adapter_checkpoint, device, ) boundary_payload = torch.load(args.boundary_checkpoint, map_location="cpu", weights_only=True) if model.boundary_head is None: raise RuntimeError("boundary head가 없습니다.") model.boundary_head.load_state_dict(boundary_payload["state_dict"]) initial_state = _selected_state06(model) family_map = _family_target_map06(base) baseline_validation = _authentic_metrics06( model, adapter, validation_rehearsal_x, validation_rehearsal_y, family_map, device=device, batch_size=args.batch_size, ) baseline_test = _authentic_metrics06( model, adapter, test_rehearsal_x, test_rehearsal_y, family_map, device=device, batch_size=args.batch_size, ) # 기존 frozen 설정에서 마지막 receptive-field block과 attention만 연다. for name, parameter in model.named_parameters(): parameter.requires_grad = name.startswith(( "trajectory_encoder.blocks.3.", "trajectory_encoder.attention.", "exact_head.", "family_head.", "boundary_head.", )) optimizer = torch.optim.AdamW([ {"params": model.boundary_head.parameters(), "lr": args.head_learning_rate}, {"params": [*model.exact_head.parameters(), *model.family_head.parameters()], "lr": args.classifier_learning_rate}, {"params": [ parameter for name, parameter in model.named_parameters() if name.startswith(("trajectory_encoder.blocks.3.", "trajectory_encoder.attention.")) ], "lr": args.shared_learning_rate}, ], weight_decay=1e-3) rehearsal_loader = DataLoader( TensorDataset(train_rehearsal_x, train_rehearsal_y), batch_size=args.batch_size, shuffle=True, generator=torch.Generator().manual_seed(args.seed), ) boundary_loader = DataLoader( TensorDataset(train_boundary_x, train_boundary_y), batch_size=args.batch_size, shuffle=True, generator=torch.Generator().manual_seed(args.seed + 7), ) history = [] winner = None minimum_exact = baseline_validation["exact_top1"] - args.maximum_exact_regression_pp / 100.0 minimum_family = baseline_validation["family_top1"] - args.maximum_family_regression_pp / 100.0 for epoch in range(1, args.epochs + 1): model.train() adapter.eval() rehearsal_iterator = iter(rehearsal_loader) boundary_iterator = iter(boundary_loader) losses = [] for _step in range(max(len(rehearsal_loader), len(boundary_loader))): try: authentic_x, authentic_y = next(rehearsal_iterator) except StopIteration: rehearsal_iterator = iter(rehearsal_loader) authentic_x, authentic_y = next(rehearsal_iterator) try: candidate_x, candidate_y = next(boundary_iterator) except StopIteration: boundary_iterator = iter(boundary_loader) candidate_x, candidate_y = next(boundary_iterator) authentic_x, authentic_y = authentic_x.to(device), authentic_y.to(device) candidate_x, candidate_y = candidate_x.to(device), candidate_y.to(device) optimizer.zero_grad(set_to_none=True) authentic_exact, authentic_family, authentic_boundary = model.forward_online_with_boundary( adapter(authentic_x), ) _candidate_exact, _candidate_family, candidate_boundary = model.forward_online_with_boundary( adapter(candidate_x), ) exact_loss = torch.nn.functional.cross_entropy(authentic_exact, authentic_y) family_loss = torch.nn.functional.cross_entropy( authentic_family, family_map[authentic_y.cpu()].to(device), ) authentic_boundary_loss = boundary_auxiliary_loss06( authentic_boundary, torch.zeros_like(authentic_boundary), ) candidate_boundary_loss = boundary_auxiliary_loss06(candidate_boundary, candidate_y) loss = ( exact_loss + 0.15 * family_loss + args.boundary_loss_weight * ( args.authentic_negative_weight * authentic_boundary_loss + candidate_boundary_loss ) ) loss.backward() torch.nn.utils.clip_grad_norm_( [parameter for parameter in model.parameters() if parameter.requires_grad], 1.0, ) optimizer.step() losses.append(float(loss.detach())) validation_authentic = _authentic_metrics06( model, adapter, validation_rehearsal_x, validation_rehearsal_y, family_map, device=device, batch_size=args.batch_size, ) validation_boundary_logits = _boundary_logits06( model, adapter, validation_boundary_x, device=device, batch_size=args.batch_size, ) threshold, _trials = _select_threshold06(validation_boundary_logits, validation_boundary_y) eligible = bool( threshold["recall_gate_passed"] and validation_authentic["exact_top1"] >= minimum_exact and validation_authentic["family_top1"] >= minimum_family ) row = { "epoch": epoch, "loss": float(np.mean(losses)), "authentic": validation_authentic, "boundary": threshold, "eligible": eligible, } history.append(row) if eligible and ( winner is None or ( threshold["f1"], validation_authentic["exact_top1"], validation_authentic["family_top1"] ) > ( winner["boundary"]["f1"], winner["authentic"]["exact_top1"], winner["authentic"]["family_top1"] ) ): winner = {**row, "state_dict": _selected_state06(model)} adopted = winner is not None if adopted: model.load_state_dict(winner.pop("state_dict"), strict=False) selected_threshold = float(winner["boundary"]["threshold"]) else: model.load_state_dict(initial_state, strict=False) selected_threshold = float(boundary_payload["threshold"]) official_authentic = _authentic_metrics06( model, adapter, test_rehearsal_x, test_rehearsal_y, family_map, device=device, batch_size=args.batch_size, ) official_boundary_logits = _boundary_logits06( model, adapter, test_boundary_x, device=device, batch_size=args.batch_size, ) official_boundary = _metrics06( official_boundary_logits, test_boundary_y, threshold=selected_threshold, ) official_deltas = { "exact_top1_pp": (official_authentic["exact_top1"] - baseline_test["exact_top1"]) * 100.0, "family_top1_pp": (official_authentic["family_top1"] - baseline_test["family_top1"]) * 100.0, } official_gate = bool( adopted and official_deltas["exact_top1_pp"] >= -args.maximum_exact_regression_pp and official_deltas["family_top1_pp"] >= -args.maximum_family_regression_pp and official_boundary["single_symbol_recall"] >= 0.95 and official_boundary["cross_boundary_recall"] >= 0.95 ) args.output.mkdir(parents=True, exist_ok=True) checkpoint_path = args.output / "boundary_joint_delta.pt" torch.save({ "schema": "aiflow-math-ink-06-p-boundary-joint-delta-v1", "state_dict": _selected_state06(model) if adopted else {}, "threshold": selected_threshold, "base_checkpoint_sha256": _file_sha25606(args.base_checkpoint), "adapter_checkpoint_sha256": _file_sha25606(args.adapter_checkpoint), "boundary_checkpoint_sha256": _file_sha25606(args.boundary_checkpoint), "training_cache_key": train_cache_key, "validation_cache_key": validation_cache_key, "test_cache_key": test_cache_key, "adopted": adopted and official_gate, "track": "P_with_obligations", "product_validation": False, }, checkpoint_path) report = { "experiment": "P-MATH-INK-06-BOUNDARY-JOINT-001", "generated_at": datetime.now(timezone.utc).isoformat(), "seed": args.seed, "device": str(device), "cuda_device": torch.cuda.get_device_name(device) if device.type == "cuda" else None, "training_contract": { "layouts": ["same_row", "superscript", "subscript", "fraction_slots", "wide_infix_sides"], "rehearsal_samples": len(train_rehearsal_y), "boundary_samples": len(train_boundary_y), "opened_modules": [ "trajectory_encoder.blocks.3", "trajectory_encoder.attention", "exact_head", "family_head", "boundary_head", ], }, "baseline_validation": baseline_validation, "winner_validation": winner, "history": history, "baseline_test": baseline_test, "official_test": { "authentic": official_authentic, "boundary": official_boundary, "deltas": official_deltas, }, "decision": { "validation_adopted": adopted, "official_gate_passed": official_gate, "release_adopted": adopted and official_gate, }, "checkpoint": str(checkpoint_path), "checkpoint_sha256": _file_sha25606(checkpoint_path), "track": "P_with_obligations", "product_validation": False, "interpretation_limit": "분리된 P 고립기호의 formula-layout 합성 proxy이며 실제 연속식 제품 gate가 아니다.", } (args.output / "report.json").write_text( json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(json.dumps({ key: value for key, value in report.items() if key != "history" }, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()