scripts/train_math_ink_06_formula_adapter.py

"""CROHME truth group에서 연속식 domain residual adapter의 구조적 회복 가능성을 진단한다."""

from __future__ import annotations

import argparse
from collections import Counter
from copy import deepcopy
from datetime import datetime, timezone
import json
from pathlib import Path
import random
import sys
from typing import Sequence

import numpy as np
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset, WeightedRandomSampler

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.skeleton_adapter06 import SkeletonTrajectoryAdapter06
from math_grid_drawer.research.trajectory_sequence import shape_family, visual_label_family
from scripts.audit_math_ink_06_case_context import _load_model06
from scripts.audit_math_ink_06_formula_context_contract import _materialize_formula_groups06
from scripts.crohme_lattice_common import writer_fit_validation
from scripts.train_crohme_segmentation_lattice_selector import _samples


def _parse_args() -> argparse.Namespace:
    """필요 변수: 제품 encoder adapter·공식 CROHME split·학습 설정. 작동 원리: R-track shadow adapter CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Train Math Ink 0.6 formula domain adapter")
    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("--profile", default="median_height_32")
    parser.add_argument("--seed", type=int, default=17)
    parser.add_argument("--epochs", type=int, default=12)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--learning-rate", type=float, default=4e-4)
    parser.add_argument("--weight-decay", type=float, default=2e-3)
    parser.add_argument("--exact-loss-weight", type=float, default=0.20)
    parser.add_argument("--context-dropout", type=float, default=0.20)
    parser.add_argument("--hidden-size", type=int, default=48)
    parser.add_argument("--patience", type=int, default=4)
    parser.add_argument("--skip-official-test", action="store_true")
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    parser.add_argument("--output", type=Path, required=True)
    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 _targets06(
    truths: Sequence[str],
    exact_labels: Sequence[str],
    family_labels: Sequence[str],
) -> tuple[Tensor, Tensor]:
    """필요 변수: truth label·model ontology. 작동 원리: exact/family target index를 만든다."""

    exact_index = {str(label): index for index, label in enumerate(exact_labels)}
    family_index = {str(label): index for index, label in enumerate(family_labels)}
    return (
        torch.tensor([exact_index[str(label)] for label in truths], dtype=torch.long),
        torch.tensor([
            family_index[shape_family(str(label))] for label in truths
        ], dtype=torch.long),
    )


def _metrics06(
    exact_logits: Tensor,
    family_logits: Tensor,
    exact_targets: Tensor,
    family_targets: Tensor,
    exact_labels: Sequence[str],
) -> dict[str, float | int]:
    """필요 변수: exact/family logit·target. 작동 원리: exact·shape-family·visual-family 지표를 같은 분모에서 계산한다."""

    exact_prediction = exact_logits.argmax(dim=1)
    family_prediction = family_logits.argmax(dim=1)
    visual_hits = sum(
        visual_label_family(str(exact_labels[int(prediction)]))
        == visual_label_family(str(exact_labels[int(truth)]))
        for prediction, truth in zip(
            exact_prediction.tolist(), exact_targets.tolist(), strict=True,
        )
    )
    return {
        "samples": len(exact_targets),
        "exact_top1": float(exact_prediction.eq(exact_targets).float().mean()),
        "family_head_top1": float(family_prediction.eq(family_targets).float().mean()),
        "visual_family_top1": visual_hits / max(len(exact_targets), 1),
    }


def _forward06(
    model,
    online_adapter,
    formula_adapter,
    features: Tensor,
    *,
    device: torch.device,
    batch_size: int,
) -> tuple[Tensor, Tensor]:
    """필요 변수: frozen main·formula adapter·feature. 작동 원리: 전체 split logit을 순서대로 CPU에 반환한다."""

    exact_rows, family_rows = [], []
    model.eval()
    online_adapter.eval()
    formula_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.classify_trajectory(
                formula_adapter(online_adapter(batch)),
            )
            exact_rows.append(exact.cpu())
            family_rows.append(family.cpu())
    return torch.cat(exact_rows), torch.cat(family_rows)


def _weighted_loader06(
    dataset: TensorDataset,
    exact_targets: Tensor,
    *,
    batch_size: int,
    seed: int,
) -> DataLoader:
    """필요 변수: 학습 tensor·exact target. 작동 원리: label 빈도 역제곱근 sampler로 대형 class 독점을 완화한다."""

    counts = torch.bincount(exact_targets).float()
    class_weight = counts.clamp_min(1.0).rsqrt()
    weights = class_weight[exact_targets]
    sampler = WeightedRandomSampler(
        weights, num_samples=len(dataset), replacement=True,
        generator=torch.Generator().manual_seed(seed),
    )
    return DataLoader(dataset, batch_size=batch_size, sampler=sampler)


def main() -> None:
    """필요 변수: writer fit/validation·official test. 작동 원리: frozen product encoder 앞 residual만 학습해 구조 상한을 측정한다."""

    args = _parse_args()
    if args.train_root.name.casefold() != "traindata" or args.test_root.name.casefold() != "testdatagt":
        raise ValueError("CROHME2012 공식 trainData/testDataGT 조합만 허용합니다.")
    if not 0.0 <= args.context_dropout <= 1.0:
        raise ValueError("context dropout은 0~1 범위여야 합니다.")
    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, online_adapter = _load_model06(base_checkpoint, args.adapter, device)
    for parameter in engine.model.parameters():
        parameter.requires_grad_(False)
    for parameter in online_adapter.parameters():
        parameter.requires_grad_(False)
    fit_samples, validation_samples = writer_fit_validation(args.train_root, args.profile)
    test_samples = [] if args.skip_official_test else _samples(args.test_root, args.profile)
    train_x, train_truths, _train_behavior = _materialize_formula_groups06(
        fit_samples, engine.labels,
    )
    validation_x, validation_truths, _validation_behavior = _materialize_formula_groups06(
        validation_samples, engine.labels,
    )
    if test_samples:
        test_x, test_truths, _test_behavior = _materialize_formula_groups06(
            test_samples, engine.labels,
        )
    else:
        test_x = torch.empty((0, 128, 19), dtype=train_x.dtype)
        test_truths = []
    train_exact, train_family = _targets06(
        train_truths, engine.labels, engine.family_labels,
    )
    validation_exact, validation_family = _targets06(
        validation_truths, engine.labels, engine.family_labels,
    )
    test_exact, test_family = (
        _targets06(test_truths, engine.labels, engine.family_labels)
        if test_truths else (
            torch.empty(0, dtype=torch.long),
            torch.empty(0, dtype=torch.long),
        )
    )
    training = TensorDataset(train_x, train_exact, train_family)
    loader = _weighted_loader06(
        training, train_exact, batch_size=args.batch_size, seed=args.seed,
    )
    formula_adapter = SkeletonTrajectoryAdapter06(hidden_size=args.hidden_size).to(device)
    optimizer = torch.optim.AdamW(
        formula_adapter.parameters(), lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=max(args.epochs, 1), eta_min=args.learning_rate * 0.1,
    )
    baseline_exact, baseline_family = _forward06(
        engine.model, online_adapter, torch.nn.Identity().to(device), validation_x,
        device=device, batch_size=args.batch_size,
    )
    baseline_validation = _metrics06(
        baseline_exact, baseline_family, validation_exact, validation_family,
        engine.labels,
    )
    best_key = (-1.0, -1.0)
    best_state = None
    best_epoch = 0
    stale = 0
    history = []
    for epoch in range(1, args.epochs + 1):
        formula_adapter.train()
        losses = []
        for features, exact_target, family_target in loader:
            features = features.to(device)
            exact_target = exact_target.to(device)
            family_target = family_target.to(device)
            if args.context_dropout:
                drop = torch.rand(len(features), device=device) < args.context_dropout
                features = features.clone()
                features[drop, :, 10:15] = 0.0
            optimizer.zero_grad(set_to_none=True)
            with torch.no_grad():
                online = online_adapter(features)
            exact, family = engine.model.classify_trajectory(
                formula_adapter(online),
            )
            loss = (
                torch.nn.functional.cross_entropy(family, family_target)
                + args.exact_loss_weight
                * torch.nn.functional.cross_entropy(exact, exact_target)
            )
            loss.backward()
            torch.nn.utils.clip_grad_norm_(formula_adapter.parameters(), 2.0)
            optimizer.step()
            losses.append(float(loss.detach()))
        scheduler.step()
        validation_logits = _forward06(
            engine.model, online_adapter, formula_adapter, validation_x,
            device=device, batch_size=args.batch_size,
        )
        validation_metrics = _metrics06(
            *validation_logits, validation_exact, validation_family, engine.labels,
        )
        row = {
            "epoch": epoch,
            "loss": sum(losses) / max(len(losses), 1),
            "validation": validation_metrics,
        }
        history.append(row)
        print(json.dumps(row, ensure_ascii=False), flush=True)
        key = (
            float(validation_metrics["family_head_top1"]),
            float(validation_metrics["visual_family_top1"]),
        )
        if key > best_key:
            best_key = key
            best_epoch = epoch
            best_state = deepcopy({
                name: value.detach().cpu()
                for name, value in formula_adapter.state_dict().items()
            })
            stale = 0
        else:
            stale += 1
        if stale >= args.patience:
            break
    if best_state is None:
        raise RuntimeError("formula adapter checkpoint가 선택되지 않았습니다.")
    formula_adapter.load_state_dict(best_state)
    selected_validation_logits = _forward06(
        engine.model, online_adapter, formula_adapter, validation_x,
        device=device, batch_size=args.batch_size,
    )
    selected_validation = _metrics06(
        *selected_validation_logits,
        validation_exact, validation_family, engine.labels,
    )
    if len(test_x):
        test_logits = _forward06(
            engine.model, online_adapter, formula_adapter, test_x,
            device=device, batch_size=args.batch_size,
        )
        official_test = _metrics06(
            *test_logits, test_exact, test_family, engine.labels,
        )
    else:
        official_test = None
    args.output.mkdir(parents=True, exist_ok=True)
    checkpoint = args.output / "formula_adapter.pt"
    torch.save({
        "schema": "aiflow-math-ink-06-formula-adapter-r-v1",
        "state_dict": best_state,
        "hidden_size": args.hidden_size,
        "base_checkpoint": str(base_checkpoint),
        "online_adapter": str(args.adapter),
        "selected_epoch": best_epoch,
        "context_dropout": args.context_dropout,
        "exact_loss_weight": args.exact_loss_weight,
        "track": "R_noncommercial_only",
        "product_validation": False,
        "distillation_allowed": False,
    }, checkpoint)
    report = {
        "experiment": "R-MATH-INK-06-FORMULA-ADAPTER-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
        ),
        "split_contract": "CROHME trainData writer fit/validation; testDataGT one-shot",
        "samples": {
            "fit": len(train_x),
            "validation": len(validation_x),
            "official_test": len(test_x),
        },
        "label_support": {
            "fit": len(Counter(train_truths)),
            "validation": len(Counter(validation_truths)),
            "official_test": len(Counter(test_truths)),
        },
        "baseline_validation": baseline_validation,
        "selected_epoch": best_epoch,
        "selected_validation": selected_validation,
        "official_test": official_test,
        "official_test_skipped": args.skip_official_test,
        "history": history,
        "checkpoint": checkpoint.name,
        "checkpoint_bytes": checkpoint.stat().st_size,
        "track": "R_noncommercial_only",
        "product_validation": False,
        "distillation_allowed": False,
    }
    (args.output / "report.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({
        "selected_epoch": best_epoch,
        "baseline_validation": baseline_validation,
        "selected_validation": selected_validation,
        "official_test": official_test,
        "product_validation": False,
    }, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()