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