scripts/export_math_ink_06_p_mobile_pair.py

"""통과한 P Formula student online과 5-output raster를 하나의 모바일 쌍으로 export한다."""

from __future__ import annotations

import argparse
import importlib.util
import json
from pathlib import Path
import sys
from typing import Any

import numpy as np
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.ink06_canonical import render_canonical_ink
from math_grid_drawer.research.ink06_export import (
    PFormulaStudentExportWrapper06,
    RasterDebugExportWrapper06,
    exported_equivalence06,
)
from math_grid_drawer.research.math_ink_06 import MathInk06Engine
from math_grid_drawer.research.p_formula_dataset06 import (
    _formula_box06,
    p_formula_symbol_ink06,
)
from math_grid_drawer.research.p_formula_gate06 import audit_p_formula_records06
from math_grid_drawer.research.skeleton_adapter06 import (
    DualModalityTrajectoryAdapter06,
    SkeletonTrajectoryAdapter06,
)
from scripts.export_math_ink_06_litert import (
    _convert_litert,
    _save_exported_program06,
    _vocabulary_sha25606,
)
from scripts.export_math_ink_06_p_formula_student import (
    _online_branch06,
    _file_sha25606,
    _read_jsonl06,
    _resolve_checkpoint_path06,
    validate_p_formula_student_artifacts06,
    validate_p_formula_student_export06,
)


MAXIMUM_MODEL_BUNDLE_BYTES06 = 25 * 1024 * 1024


def paired_p_representatives06(
    records: list[dict[str, Any]],
    *,
    maximum_samples: int,
) -> tuple[list[tuple[torch.Tensor, ...]], list[tuple[torch.Tensor, ...]]]:
    """필요 변수: P test formula·상한. 작동 원리: 같은 symbol에서 online 128×19와 raster 128×128을 함께 만든다."""

    if maximum_samples <= 0:
        raise ValueError("대표 입력 상한은 양수여야 합니다.")
    online: list[tuple[torch.Tensor, ...]] = []
    raster: list[tuple[torch.Tensor, ...]] = []
    for record in records:
        if str(record.get("split") or "") != "test":
            continue
        formula_box = _formula_box06(record)
        for symbol in record["symbols"]:
            ink = p_formula_symbol_ink06(symbol, formula_box=formula_box)
            image = np.asarray(render_canonical_ink(ink), dtype=np.float32)
            online.append((torch.from_numpy(ink.features).unsqueeze(0),))
            raster.append((
                torch.from_numpy(1.0 - image / 255.0).unsqueeze(0).unsqueeze(0),
            ))
            if len(online) >= maximum_samples:
                return online, raster
    if not online:
        raise ValueError("P Formula test representative가 없습니다.")
    return online, raster


def _raster_branch06(adapter: torch.nn.Module) -> torch.nn.Module:
    """필요 변수: single/dual adapter. 작동 원리: image virtual stroke에 대응하는 raster branch만 반환한다."""

    return adapter.raster if isinstance(adapter, DualModalityTrajectoryAdapter06) else adapter


def main() -> None:
    """필요 변수: 통과 student·동일 P corpus·출력. 작동 원리: 같은 lineage의 online/raster graph를 export한다."""

    parser = argparse.ArgumentParser(description="Export Math Ink 0.6 P mobile pair")
    parser.add_argument("--student-checkpoint", type=Path, required=True)
    parser.add_argument("--base-checkpoint", type=Path)
    parser.add_argument("--adapter-checkpoint", type=Path)
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--maximum-representative-samples", type=int, default=256)
    parser.add_argument("--convert-litert", action="store_true")
    args = parser.parse_args()
    payload = torch.load(args.student_checkpoint, map_location="cpu", weights_only=False)
    data_sha256 = _file_sha25606(args.data)
    validate_p_formula_student_export06(payload, data_sha256=data_sha256)
    parent = args.student_checkpoint.parent
    base = args.base_checkpoint or _resolve_checkpoint_path06(
        payload["student_base_checkpoint"],
        parent=parent,
    )
    adapter_checkpoint = args.adapter_checkpoint or _resolve_checkpoint_path06(
        payload["student_online_adapter"],
        parent=parent,
    )
    validate_p_formula_student_artifacts06(
        payload,
        base_checkpoint=base,
        online_adapter=adapter_checkpoint,
    )
    engine = MathInk06Engine(base, adapter_checkpoint=adapter_checkpoint)
    formula_adapter = SkeletonTrajectoryAdapter06(
        hidden_size=int(payload["hidden_size"]),
    )
    formula_adapter.load_state_dict(payload["state_dict"])
    online_wrapper = PFormulaStudentExportWrapper06(
        engine.model,
        _online_branch06(engine.composite_adapter),
        formula_adapter,
        family_weight=engine.online_family_fusion_weight,
        exact_family_index=engine.exact_family_index,
    ).eval()
    fusion = engine.raster_fusion
    if any(float(fusion[key]) != 0.0 for key in ("family_weight", "geometry_weight", "symmetry_weight")):
        raise ValueError("Mobile raster export는 학습 graph 밖 auxiliary fusion을 허용하지 않습니다.")
    raster_wrapper = RasterDebugExportWrapper06(
        engine.model,
        adapter=_raster_branch06(engine.composite_adapter),
        fusion_mode=str(fusion["mode"]),
        score_weight=float(fusion["score_weight"]),
    ).eval()
    records = _read_jsonl06(args.data)
    audit = audit_p_formula_records06(records)
    if not audit["eligible_for_product_evaluation"]:
        raise ValueError("P Formula corpus가 product preflight를 통과하지 못했습니다.")
    online_inputs, raster_inputs = paired_p_representatives06(
        records,
        maximum_samples=args.maximum_representative_samples,
    )
    online_export = torch.export.export(
        online_wrapper,
        online_inputs[0],
        strict=True,
    )
    raster_export = torch.export.export(
        raster_wrapper,
        raster_inputs[0],
        strict=True,
    )
    equivalence = {
        "online": exported_equivalence06(
            online_wrapper,
            online_export,
            online_inputs,
        ),
        "raster": exported_equivalence06(
            raster_wrapper,
            raster_export,
            raster_inputs,
        ),
    }
    args.output.mkdir(parents=True, exist_ok=True)
    online_path = args.output / "p_formula_online.pt2"
    raster_path = args.output / "raster_debug5.pt2"
    _save_exported_program06(online_export, online_path)
    _save_exported_program06(raster_export, raster_path)
    total_bytes = online_path.stat().st_size + raster_path.stat().st_size
    size_gate = total_bytes <= MAXIMUM_MODEL_BUNDLE_BYTES06
    labels = tuple(str(label) for label in engine.labels)
    report: dict[str, Any] = {
        "schema": "aiflow-math-ink-06-p-mobile-pair-export-v1",
        "model_version": f"{engine.model_version}+p-formula-student",
        "student_checkpoint": str(args.student_checkpoint),
        "data_sha256": data_sha256,
        "teacher_seeds": [17, 31, 47],
        "teacher_weights_embedded": False,
        "exact_label_count": len(labels),
        "vocabulary_sha256": _vocabulary_sha25606(list(labels)),
        "raster_output_count": 5,
        "representative_samples": len(online_inputs),
        "preflight": audit,
        "torch_version": torch.__version__,
        "torch_export": {
            "online": {
                **equivalence["online"],
                "path": online_path.name,
                "bytes": online_path.stat().st_size,
            },
            "raster": {
                **equivalence["raster"],
                "path": raster_path.name,
                "bytes": raster_path.stat().st_size,
            },
            "total_bytes": total_bytes,
            "maximum_bundle_bytes": MAXIMUM_MODEL_BUNDLE_BYTES06,
            "size_gate_passed": size_gate,
        },
        "torch_export_gate_passed": bool(
            size_gate
            and equivalence["online"]["gate_passed"]
            and equivalence["raster"]["gate_passed"]
        ),
        "litert_package_available": importlib.util.find_spec("litert_torch") is not None,
        "litert": {
            "online": {"converted": False, "reason": "conversion_not_requested"},
            "raster": {"converted": False, "reason": "conversion_not_requested"},
        },
        "product_validation": False,
    }
    if args.convert_litert:
        if not report["litert_package_available"]:
            for branch in ("online", "raster"):
                report["litert"][branch] = {
                    "converted": False,
                    "reason": "litert_torch_not_installed",
                }
        else:
            report["litert"] = {
                "online": _convert_litert(
                    online_wrapper,
                    online_inputs,
                    args.output / "p_formula_online.tflite",
                ),
                "raster": _convert_litert(
                    raster_wrapper,
                    raster_inputs,
                    args.output / "raster_debug5.tflite",
                ),
            }
    report["next_gate"] = (
        "package exact model pair, raster release validation, Android low/mid/high"
    )
    (args.output / "export_manifest.json").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()