scripts/build_math_ink_06_litert_colab_bundle.py

"""AIFlow Math Ink 0.6 LiteRT 변환에 필요한 최소 Colab ZIP을 결정적으로 만든다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
from hashlib import sha256
import json
from pathlib import Path, PurePosixPath
from zipfile import ZIP_DEFLATED, ZIP_STORED, ZipFile, ZipInfo


PROJECT_ROOT = Path(__file__).parents[1]
MANIFEST_NAME_06 = "LITERT_COLAB_BUNDLE_MANIFEST.json"


def _sha256_bytes06(value: bytes) -> str:
    """필요 변수: file bytes. 작동 원리: bundle entry와 manifest 검증용 SHA-256을 반환한다."""

    return sha256(value).hexdigest()


def _bundle_files06(
    base_checkpoint: Path, adapter_checkpoint: Path, representative_inputs: Path,
) -> dict[str, Path]:
    """필요 변수: seed-17 composite artifact. 작동 원리: converter가 import하는 닫힌 최소 파일 집합을 매핑한다."""

    return {
        "pyproject.toml": PROJECT_ROOT / "pyproject.toml",
        "src/math_grid_drawer/__init__.py": PROJECT_ROOT / "research/colab/litert_bundle_package_init.py",
        "src/math_grid_drawer/research/__init__.py": PROJECT_ROOT / "src/math_grid_drawer/research/__init__.py",
        "src/math_grid_drawer/research/ink06_canonical.py": PROJECT_ROOT / "src/math_grid_drawer/research/ink06_canonical.py",
        "src/math_grid_drawer/research/ink06_export.py": PROJECT_ROOT / "src/math_grid_drawer/research/ink06_export.py",
        "src/math_grid_drawer/research/math_ink_06.py": PROJECT_ROOT / "src/math_grid_drawer/research/math_ink_06.py",
        "src/math_grid_drawer/research/raster_skeleton06.py": PROJECT_ROOT / "src/math_grid_drawer/research/raster_skeleton06.py",
        "src/math_grid_drawer/research/skeleton_adapter06.py": PROJECT_ROOT / "src/math_grid_drawer/research/skeleton_adapter06.py",
        "src/math_grid_drawer/research/trajectory_sequence.py": PROJECT_ROOT / "src/math_grid_drawer/research/trajectory_sequence.py",
        "scripts/export_math_ink_06_litert.py": PROJECT_ROOT / "scripts/export_math_ink_06_litert.py",
        "artifacts/base_378.pt": base_checkpoint,
        "artifacts/online_adapter.pt": adapter_checkpoint,
        "artifacts/representative_inputs.pt": representative_inputs,
    }


def build_litert_colab_bundle06(
    base_checkpoint: Path,
    adapter_checkpoint: Path,
    representative_inputs: Path,
    output: Path,
) -> dict:
    """필요 변수: composite checkpoint·대표 입력·ZIP 출력. 작동 원리: 경로·크기·해시가 고정된 Linux 변환 bundle을 만든다."""

    files = _bundle_files06(base_checkpoint, adapter_checkpoint, representative_inputs)
    missing = [name for name, path in files.items() if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"LiteRT bundle 필수 파일이 없습니다: {missing}")
    entries = []
    payloads: dict[str, bytes] = {}
    for name, path in files.items():
        payload = path.read_bytes()
        payloads[name] = payload
        entries.append({
            "path": name, "bytes": len(payload), "sha256": _sha256_bytes06(payload),
        })
    manifest = {
        "schema": "aiflow-math-ink-06-litert-colab-bundle-v1",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "track": "R_public_conversion_only",
        "p_student_included": False,
        "product_bundle": False,
        "seed": 17,
        "litert_torch_version": "0.9.1",
        "representative_samples": 76,
        "raster_output_contract": {
            "outputs": [
                {"index": 0, "name": "exact_logits", "shape": [1, 378]},
                {"index": 1, "name": "coordinates", "shape": [1, 4, 128, 2]},
                {"index": 2, "name": "state_logits", "shape": [1, 4, 128, 3]},
                {"index": 3, "name": "progress", "shape": [1, 4, 128]},
                {"index": 4, "name": "hypothesis_scores", "shape": [1, 4]},
            ],
            "direct_raster_label_shortcut": False,
        },
        "files": entries,
        "product_validation": False,
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(output.suffix + ".part")
    with ZipFile(temporary, "w", allowZip64=True) as bundle:
        for name, payload in payloads.items():
            compression = ZIP_STORED if PurePosixPath(name).suffix in {".pt"} else ZIP_DEFLATED
            info = ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0))
            info.compress_type = compression
            info.external_attr = 0o644 << 16
            bundle.writestr(info, payload)
        info = ZipInfo(MANIFEST_NAME_06, date_time=(1980, 1, 1, 0, 0, 0))
        info.compress_type = ZIP_DEFLATED
        info.external_attr = 0o644 << 16
        bundle.writestr(
            info, json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8") + b"\n",
        )
    temporary.replace(output)
    manifest["bundle"] = str(output)
    manifest["bundle_bytes"] = output.stat().st_size
    manifest["bundle_sha256"] = sha256(output.read_bytes()).hexdigest()
    return manifest


def verify_litert_colab_bundle06(bundle_path: Path) -> dict:
    """필요 변수: 생성 ZIP. 작동 원리: 추출 없이 모든 entry의 안전 경로·크기·SHA-256을 재검증한다."""

    failures = []
    with ZipFile(bundle_path) as bundle:
        manifest = json.loads(bundle.read(MANIFEST_NAME_06).decode("utf-8"))
        if (
            manifest.get("track") != "R_public_conversion_only"
            or manifest.get("p_student_included") is not False
            or manifest.get("product_bundle") is not False
        ):
            failures.append({"path": MANIFEST_NAME_06, "reason": "research_track_contract"})
        contract = manifest.get("raster_output_contract") or {}
        output_names = [
            str(row.get("name") or "")
            for row in contract.get("outputs", [])
        ]
        if output_names != [
            "exact_logits",
            "coordinates",
            "state_logits",
            "progress",
            "hypothesis_scores",
        ]:
            failures.append({"path": MANIFEST_NAME_06, "reason": "raster_output_contract"})
        if contract.get("direct_raster_label_shortcut") is not False:
            failures.append({"path": MANIFEST_NAME_06, "reason": "raster_shortcut_contract"})
        names = set(bundle.namelist())
        for row in manifest["files"]:
            name = str(row["path"])
            path = PurePosixPath(name)
            if path.is_absolute() or ".." in path.parts:
                failures.append({"path": name, "reason": "unsafe_path"})
            elif name not in names:
                failures.append({"path": name, "reason": "missing"})
            else:
                payload = bundle.read(name)
                if len(payload) != int(row["bytes"]):
                    failures.append({"path": name, "reason": "bytes"})
                elif _sha256_bytes06(payload) != str(row["sha256"]):
                    failures.append({"path": name, "reason": "sha256"})
    return {
        "schema": manifest["schema"], "files": len(manifest["files"]),
        "failures": failures, "passed": not failures, "product_validation": False,
    }


def main() -> None:
    """필요 변수: CLI artifact 경로. 작동 원리: bundle 생성 직후 archive 자체 검증까지 수행한다."""

    parser = argparse.ArgumentParser(description="Build Math Ink 0.6 LiteRT Colab bundle")
    parser.add_argument("--base-checkpoint", type=Path, required=True)
    parser.add_argument("--adapter-checkpoint", type=Path, required=True)
    parser.add_argument("--representative-inputs", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = build_litert_colab_bundle06(
        args.base_checkpoint, args.adapter_checkpoint, args.representative_inputs, args.output,
    )
    report["verification"] = verify_litert_colab_bundle06(args.output)
    if not report["verification"]["passed"]:
        raise ValueError(f"LiteRT Colab bundle 검증 실패: {report['verification']['failures']}")
    report_path = args.output.with_suffix(args.output.suffix + ".manifest.json")
    report_path.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()