File size: 4,932 Bytes
8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 56ee0e9 8a48617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Command-line entry point shared by the Kai and Lex Core ML toolkits."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from huggingface_hub import snapshot_download
from typed_coreml import (
KINDS,
compare_native,
export_coreml,
load_native,
verify_coreml,
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--toolkit", required=True, type=Path)
parser.add_argument("--native-dir", type=Path)
parser.add_argument("--output-dir", type=Path)
parser.add_argument("--kind", choices=KINDS)
parser.add_argument("--tokens", type=int, default=128)
parser.add_argument("--candidates", type=int)
parser.add_argument("--fixture-index", type=int, default=0)
parser.add_argument("--flexible", action="store_true")
parser.add_argument("--marker-map", action="store_true", help="use an FP16 one-hot marker map instead of indices")
parser.add_argument("--verify-only", action="store_true")
parser.add_argument("--verify-limit", type=int, help="maximum pinned real rows to verify per type")
parser.add_argument("--verify-package", type=Path)
parser.add_argument("--verify-packages-dir", type=Path)
args = parser.parse_args()
if args.marker_map and args.flexible:
parser.error("--marker-map currently requires a fixed token and candidate shape")
if args.verify_limit is not None and args.verify_limit < 1:
parser.error("--verify-limit must be positive")
if args.verify_package is not None and args.kind is None:
parser.error("--verify-package requires --kind")
if args.verify_package is not None and args.verify_packages_dir is not None:
parser.error("Choose one package verification mode")
toolkit = args.toolkit.resolve()
lock = json.loads((toolkit / "assets.lock.json").read_text())
native_dir = args.native_dir
if native_dir is None:
snapshot = snapshot_download(
repo_id=lock["source_repo"],
revision=lock["source_revision"],
allow_patterns=["native/**"],
)
native_dir = Path(snapshot) / "native"
model, collator, _ = load_native(native_dir)
records = []
for fixture in ("upstream-decisions.jsonl", "upstream-system-one-rows.jsonl"):
records.extend(json.loads(line) for line in (toolkit / fixture).read_text().splitlines() if line)
kinds = [args.kind] if args.kind else list(KINDS)
failed_packages = []
for kind in kinds:
selected = [record for record in records if record["question"]["type"].lower() == kind]
if not selected:
raise ValueError(f"Missing upstream {kind} fixture")
if args.verify_only:
print(json.dumps(compare_native(
model, collator, selected, kind, tokens=args.tokens, candidates=args.candidates,
marker_map=args.marker_map,
)))
continue
package = args.verify_package
if args.verify_packages_dir is not None:
package = args.verify_packages_dir / (kind + ".mlpackage")
if package is not None:
report = verify_coreml(
model, collator, selected[:args.verify_limit], kind, package
)
output = package.parent / (kind + "-validation.json")
output.write_text(json.dumps(report, indent=2) + "\n")
metadata_path = package.with_suffix(".json")
metadata = json.loads(metadata_path.read_text())
passed = report["choice_agreement"] and report["max_abs_probability_error"] <= 0.02
metadata["validated"] = passed
metadata["validation_report"] = output.name
metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")
print(json.dumps(report))
if not passed:
failed_packages.append(kind)
continue
if args.output_dir is None:
raise ValueError("--output-dir is required for conversion")
if args.fixture_index < 0 or args.fixture_index >= len(selected):
raise ValueError("--fixture-index is outside the upstream fixture set")
suffix = "-flex" if args.flexible else "-marker-map" if args.marker_map else ""
report = export_coreml(
model,
collator,
[selected[args.fixture_index]],
kind,
args.output_dir / (kind + suffix + ".mlpackage"),
source_repo=lock["source_repo"],
source_revision=lock["source_revision"],
tokens=args.tokens,
candidates=args.candidates,
flexible=args.flexible,
marker_map=args.marker_map,
)
print(json.dumps(report))
if failed_packages:
raise SystemExit(f"Core ML native-parity gate failed for: {', '.join(failed_packages)}")
if __name__ == "__main__":
main()
|