scripts/train_math_ink_06_p_boundary_joint.py

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