"""Independent stage commands; user starts every real experiment manually.""" import argparse import json import math from pathlib import Path import torch from vimeml.deployment.bundle import BundleLM, ROOT, export_bundle def main(argv=None): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--threads", type=int, default=4) sub = parser.add_subparsers(dest="command", required=True) export = sub.add_parser("export", help="Windows: inference-only bundle; never alters frozen inputs.") export.add_argument("--checkpoint", type=Path, default=ROOT / "artifacts/models/tiny-ja-v1/best.pt") export.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/ja-unigram-16k-v1/tokenizer.model") export.add_argument("--token-manifest", type=Path, default=ROOT / "artifacts/token-data/corpus-v1-16k/manifest.json") export.add_argument("--output", type=Path, required=True) ref = sub.add_parser("reference", help="Windows CPU FP32 portable tensor/scoring fixtures.") convert = sub.add_parser("convert", help="Mac: uncompressed FP16 ML Program.") convert.add_argument("--target", type=int, choices=(16, 18), default=18) compress = sub.add_parser("compress", help="Mac: weight-only post-training compression.") compress.add_argument("--source", type=Path, required=True) compress.add_argument("--fp16-alignment", type=Path, required=True) compress.add_argument("--method", choices=("palette", "linear"), default="palette") compress.add_argument("--bits", type=int, choices=(4, 8), default=4) compress.add_argument("--group-size", type=int, choices=(0, 8, 16, 32), default=16, help="0=per-tensor palette, positive=grouped channel (iOS18).") compress.add_argument("--block-size", type=int, choices=(16, 32, 64, 128), default=32) alignment = sub.add_parser("validate", help="Mac: compare exact Windows FP32 fixtures.") alignment.add_argument("--reference", type=Path, required=True) alignment.add_argument("--atol", type=float, default=.1) alignment.add_argument("--rtol", type=float, default=.01) alignment.add_argument("--score-atol", type=float, default=.2) evaluate = sub.add_parser("evaluate", help="Fresh dev/AJIMEE scores; LM-only, no lambda.") evaluate.add_argument("--benchmark", type=Path, required=True) evaluate.add_argument("--role", choices=("dev", "ajimee"), required=True) evaluate.add_argument("--baseline", type=Path, help="Existing frozen FP32 or new deployment scores directory.") phrases = sub.add_parser("phrases", help="Same fixed 20-prefix demo suite; preserve raw outputs.") timing = sub.add_parser("timing", help="Mac warm load/reranking/beam latency; no device claims.") timing.add_argument("--benchmark", type=Path, required=True) timing.add_argument("--repeats", type=int, default=5) timing.add_argument("--warmup", type=int, default=1) for command in (phrases, timing): command.add_argument("--prompts", type=Path, default=ROOT / "configs/phrase-demo-prompts.json") phrases.add_argument("--mode", choices=("beam", "sample"), default="beam") for command in (ref, convert, alignment, evaluate, phrases, timing): command.add_argument("--bundle", type=Path, required=True) for command in (alignment, evaluate, phrases, timing): command.add_argument("--model", type=Path, help="Core ML experiment directory; omit for bundle FP32.") command.add_argument("--compute-units", choices=("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"), default="CPU_ONLY") for command in (ref, convert, compress, alignment, evaluate, phrases, timing): command.add_argument("--output", type=Path, required=True) args = parser.parse_args(argv) if args.threads < 1: parser.error("threads must be positive.") if args.command == "validate" and any(not math.isfinite(value) or value < 0 for value in (args.atol, args.rtol, args.score_atol)): parser.error("tolerances must be finite and nonnegative.") if args.command == "timing" and (args.repeats < 1 or args.warmup < 0): parser.error("repeats must be positive and warmup nonnegative.") if args.output.exists(): parser.error("Output exists; use a new versioned directory (no overwrite option).") torch.set_num_threads(args.threads) if args.command == "export": export_bundle(args.checkpoint, args.tokenizer, args.token_manifest, args.output) elif args.command == "convert": from vimeml.deployment.coreml import convert as run run(args.bundle, args.output, args.target) elif args.command == "compress": from vimeml.deployment.coreml import compress as run run(args.source, args.output, args.method, args.bits, args.group_size, args.block_size, args.fp16_alignment) else: if getattr(args, "model", None): from vimeml.deployment.coreml import CoreMLLM lm = CoreMLLM(args.bundle, args.model, args.compute_units) else: lm = BundleLM(args.bundle) from vimeml.deployment import validation if args.command == "reference": validation.reference(lm, args.output) elif args.command == "validate": report = validation.validate(lm, args.reference, args.output, args.atol, args.rtol, args.score_atol) print(json.dumps({key: report[key] for key in ("passed", "logits", "invariants")}, indent=2)) if not report["passed"]: raise SystemExit("Alignment failed; detailed report preserved. Inspect it before proceeding.") elif args.command == "evaluate": report = validation.evaluate(lm, args.benchmark, args.role, args.output, args.baseline) print(json.dumps(report["metrics"]["all"], ensure_ascii=False, indent=2)) elif args.command == "phrases": from vimeml.tools.phrase_demo import PhraseDemo, run_suite from vimeml.deployment.bundle import read_json run_suite(PhraseDemo(lm, read_json(args.prompts)), args.prompts, args.output, max_tokens=8, mode=args.mode) elif args.command == "timing": validation.timing(lm, args.benchmark, args.prompts, args.output, args.repeats, args.warmup) print(f"Output: {args.output.resolve()}")