| """`minijev` command line. |
| |
| minijev hardware |
| minijev train --config configs/snake-small.yaml [--resume RUN_DIR] [--set key=value ...] |
| minijev evaluate --checkpoint RUN/checkpoints/final.pt [--games 200] |
| minijev benchmark --checkpoint RUN/checkpoints/final.pt [--onnx] |
| minijev export --checkpoint RUN/checkpoints/final.pt --format onnx |
| minijev estimate --config configs/vast/vast-small.yaml --gpu-hourly-cost 0.20 |
| minijev compare RUN_DIR [RUN_DIR ...] |
| minijev runs annotate RUN_DIR --actual-cost 0.42 |
| minijev export-hf --checkpoint X.pt --out hf/name (checkpoint may then be a folder or a hub repo id) |
| minijev relabel --checkpoint X.pt --out Y.pt --map UP="move up" ... [--set task.choice_style=semantic] |
| minijev ask --checkpoint X.pt --state hotdog --question "Is this food?" --choices yes no |
| minijev bench-jsonl --checkpoint X.pt --data data/general_qa/benchmark.jsonl |
| minijev snake generate|baselines|demo|play|dagger|rl ... |
| minijev sentiment prepare ... |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
|
|
| def _cfg(args): |
| from minijev.config import load_config |
| return load_config(args.config, getattr(args, "set", None) or []) |
|
|
|
|
| |
| def cmd_hardware(args): |
| import torch |
|
|
| from minijev.device import resolve_runtime |
| prof = resolve_runtime({"batch_size": args.batch_size}, args.params) |
| print(json.dumps({"torch": torch.__version__, **prof.to_dict()}, indent=2)) |
|
|
|
|
| def cmd_train(args): |
| from minijev.training.supervised import train |
| cfg = _cfg(args) |
| run_dir = args.resume or args.run_dir |
| s = train(cfg, run_dir=run_dir, resume=bool(args.resume), max_steps=args.max_steps, max_minutes=args.max_minutes) |
| print(json.dumps({k: v for k, v in s.items() if not k.startswith("test_")}, indent=2, default=str)) |
| if cfg["task"]["name"] == "snake" and s.get("final_checkpoint") and not args.no_games: |
| _snake_eval(s["final_checkpoint"], cfg, Path(s["run_dir"])) |
|
|
|
|
| def _snake_eval(ckpt, cfg, run_dir: Path, games: int | None = None): |
| from environments.snake.evaluate import model_policy, play_games |
| from minijev.inference.runtime import MiniJEV |
| dec = MiniJEV.load(ckpt, threads=4) |
| env = cfg.get("env", {}) |
| n = games or cfg.get("eval", {}).get("games", 200) |
| res = play_games(model_policy(dec), n, env, parallel=cfg.get("eval", {}).get("parallel", 64)) |
| print(f"[games] {res['games']} games on {res['board']}: avg {res['avg_score']:.1f} median {res['median_score']:.0f} " |
| f"max {res['max_score']} / {res['max_possible_score']} | win {res['win_rate']:.2%} " |
| f"collision {res['death_rate']:.2%} starvation {res['starvation_rate']:.2%} | " |
| f"fatal-move rate {res['fatal_move_rate']:.4%}") |
| (run_dir / "games.json").write_text(json.dumps(res, indent=2)) |
| return res |
|
|
|
|
| def cmd_evaluate(args): |
| import numpy as np |
|
|
| from minijev.calibration.metrics import calibration_report |
| from minijev.datasets.batching import BatchEncoder |
| from minijev.datasets.decision_data import DecisionData |
| from minijev.inference.runtime import MiniJEV |
| dec = MiniJEV.load(args.checkpoint, threads=args.threads) |
| import torch |
| cfg = torch.load(args.checkpoint, weights_only=False)["config"] if args.checkpoint.endswith(".pt") else None |
| data_dir = Path(args.data or (cfg["data"]["dir"] if cfg else "")) |
| out = {"checkpoint": args.checkpoint} |
| split = data_dir / f"{args.split}.npz" |
| if split.exists(): |
| data = DecisionData.load(split) |
| enc = BatchEncoder(data, dec.task, dec.max_question_len, dec.max_choice_len) |
| logits = [] |
| for i in range(0, len(data), 1024): |
| b = enc(np.arange(i, min(len(data), i + 1024))) |
| logits.append(dec._run(next(v for k, v in b.items() if k.startswith("state_")), |
| b["question_ids"], b["choice_ids"])[1]) |
| acceptable = data.soft > 0 if data.soft is not None else None |
| rep = calibration_report(np.concatenate(logits), data.target, acceptable=acceptable) |
| out["calibration"] = {k: v for k, v in rep.items() if k != "reliability"} |
| out["reliability"] = rep["reliability"] |
| print(f"[{args.split}] n={rep['n']} acc {rep['accuracy']:.4f} ECE {rep['ece']:.4f} " |
| f"NLL {rep['nll']:.4f} Brier {rep['brier']:.4f}" |
| + (f" acceptable-acc {rep['acceptable_accuracy']:.4f}" if acceptable is not None else "")) |
| if dec.task.name == "snake" and args.games: |
| from environments.snake.evaluate import model_policy, play_games |
| env = (cfg or {}).get("env", {}) |
| if args.board_size: |
| env = {**env, "board_size": args.board_size} |
| res = play_games(model_policy(dec), args.games, env) |
| out["games"] = res |
| print(f"[games] avg {res['avg_score']:.2f} median {res['median_score']} max {res['max_score']} " |
| f"win {res['win_rate']:.2%} collisions {res['death_rate']:.2%} fatal-move {res['fatal_move_rate']:.4%}") |
| if args.out: |
| Path(args.out).write_text(json.dumps(out, indent=2)) |
|
|
|
|
| def cmd_benchmark(args): |
| from minijev.benchmark.perf import benchmark_decider, format_perf |
| paths = [args.checkpoint] |
| if args.onnx: |
| p = Path(args.checkpoint) |
| for cand in (p.with_suffix(".onnx"), p.with_name(p.stem + ".int8.onnx")): |
| if cand.exists(): |
| paths.append(str(cand)) |
| results = [] |
| for path in paths: |
| r = benchmark_decider(path, threads=args.threads, quick=args.quick) |
| print(format_perf(r)) |
| results.append(r) |
| if args.out: |
| Path(args.out).write_text(json.dumps(results, indent=2)) |
|
|
|
|
| def cmd_export(args): |
| from minijev.export.onnx_export import export_onnx |
| if args.format != "onnx": |
| sys.exit("only --format onnx is supported") |
| r = export_onnx(args.checkpoint, args.out, quantize=not args.no_quantize) |
| print(json.dumps({k: v for k, v in r.items() if k != "meta"}, indent=2)) |
|
|
|
|
| def cmd_estimate(args): |
| from minijev.estimate import estimate |
| estimate(_cfg(args), args.gpu_hourly_cost, throughput=args.throughput, bench_seconds=args.bench_seconds, |
| gpu_name=args.gpu_name, budget=args.budget) |
|
|
|
|
| def cmd_relabel(args): |
| from minijev.config import _parse_value |
| from minijev.training.relabel import relabel_checkpoint |
| mapping = dict(m.split("=", 1) for m in args.map) |
| updates = {k: _parse_value(v) for k, v in (kv.split("=", 1) for kv in (args.set or []))} |
| print(json.dumps(relabel_checkpoint(args.checkpoint, args.out, mapping, updates, steps=args.steps), indent=2)) |
|
|
|
|
| def cmd_export_hf(args): |
| from minijev.hub import save_pretrained |
| print(save_pretrained(args.checkpoint, args.out)) |
|
|
|
|
| def cmd_ask(args): |
| from minijev.inference.runtime import MiniJEV |
| dec = MiniJEV.load(args.checkpoint, threads=args.threads, **({"task": args.task} if args.task else {})) |
| d = dec.decide(args.state, args.question, args.choices) |
| print(json.dumps(d.to_dict(), ensure_ascii=False, indent=2)) |
|
|
|
|
| def cmd_bench_jsonl(args): |
| """Accuracy + calibration of a model on a JSONL decision file, row by row.""" |
| import numpy as np |
|
|
| from minijev.calibration.metrics import calibration_report |
| from minijev.datasets.jsonl import read_jsonl |
| from minijev.inference.runtime import MiniJEV |
| dec = MiniJEV.load(args.checkpoint, threads=args.threads, **({"task": args.task} if args.task else {})) |
| rows = read_jsonl(args.data) |
| ok, logits, targets, out = 0, [], [], [] |
| for r in rows: |
| d = dec.decide(r.get("state") or "", r["question"], r["choices"]) |
| tgt = r["target"] if isinstance(r["target"], str) else r["choices"][r["target"]] |
| hit = d.choice == tgt |
| ok += hit |
| n = max(len(x["choices"]) for x in rows) |
| logits.append(np.pad(np.array(d.logits), (0, n - len(d.logits)), constant_values=-1e9)) |
| targets.append(r["choices"].index(tgt)) |
| out.append({**r, "pred": d.choice, "confidence": round(d.confidence, 4), "correct": bool(hit)}) |
| if args.verbose: |
| print(f"{'OK ' if hit else 'XX '} {d.confidence:.2f} {r.get('state') or '-':>10} | {r['question']} " |
| f"-> {d.choice} (gold {tgt})") |
| rep = calibration_report(np.stack(logits), np.array(targets)) |
| print(f"[{args.data}] {ok}/{len(rows)} = {ok / len(rows):.3f} | ECE {rep['ece']:.4f} NLL {rep['nll']:.4f} " |
| f"Brier {rep['brier']:.4f}") |
| if args.out: |
| Path(args.out).write_text(json.dumps({"accuracy": ok / len(rows), "ece": rep["ece"], "nll": rep["nll"], |
| "brier": rep["brier"], "rows": out}, indent=2, ensure_ascii=False)) |
|
|
|
|
| def cmd_compare(args): |
| from minijev.report import compare_runs |
| print(compare_runs(args.runs, out=args.out)) |
|
|
|
|
| def cmd_runs_annotate(args): |
| p = Path(args.run) / "summary.json" |
| s = json.loads(p.read_text()) |
| s.setdefault("cost", {}) |
| if args.actual_cost is not None: |
| s["cost"]["actual_cost"] = args.actual_cost |
| if args.gpu: |
| s["cost"]["gpu"] = args.gpu |
| if args.note: |
| s["note"] = args.note |
| p.write_text(json.dumps(s, indent=2)) |
| print(json.dumps(s.get("cost"), indent=2)) |
|
|
|
|
| |
| def cmd_snake_generate(args): |
| from environments.snake.dataset import generate_splits |
| sizes = [int(x) for x in str(args.board_size).split(",")] |
| out = args.out or f"data/snake{'-'.join(map(str, sizes))}-{_short(args.samples)}" |
| s = generate_splits(out, args.samples, seed=args.seed, board_sizes=sizes, keep_prob=args.keep_prob, |
| explore_eps=args.explore_eps, workers=args.workers) |
| if args.jsonl: |
| from environments.snake.dataset import export_jsonl |
| from minijev.datasets.decision_data import DecisionData |
| for split in ("train", "val", "test"): |
| export_jsonl(DecisionData.load(Path(out) / f"{split}.npz"), Path(out) / f"{split}.jsonl") |
| print(json.dumps(s, indent=2)) |
| print(f"-> {out}") |
|
|
|
|
| def _short(n: int) -> str: |
| return f"{n // 1_000_000}m" if n >= 1_000_000 and n % 1_000_000 == 0 else (f"{n // 1000}k" if n >= 1000 else str(n)) |
|
|
|
|
| def cmd_snake_baselines(args): |
| from environments.snake.evaluate import baseline_suite |
| res = baseline_suite(args.games, {"board_size": args.board_size}, seed=args.seed) |
| print(f"{'agent':10} {'avg':>7} {'median':>7} {'max':>5} {'win':>6} {'collide':>8} {'starve':>7} {'fatal%':>8}") |
| for name, r in res.items(): |
| print(f"{name:10} {r['avg_score']:7.1f} {r['median_score']:7.0f} {r['max_score']:5d} {r['win_rate']:6.1%} " |
| f"{r['death_rate']:8.1%} {r['starvation_rate']:7.1%} {r['fatal_move_rate']:8.3%}") |
| if args.out: |
| Path(args.out).write_text(json.dumps(res, indent=2)) |
|
|
|
|
| def cmd_snake_demo(args, human=False): |
| from environments.snake.demo.server import serve |
| serve(model=args.model, board_size=args.board_size, port=args.port, agent="human" if human else args.agent, |
| open_browser=not args.no_browser) |
|
|
|
|
| def cmd_snake_dagger(args): |
| from environments.snake.dagger import run_dagger |
| run_dagger(_cfg(args), args.checkpoint, iterations=args.iterations, samples=args.samples, run_dir=args.run_dir) |
|
|
|
|
| def cmd_snake_rl(args): |
| from environments.snake.rl import train_ppo |
| train_ppo(_cfg(args), args.checkpoint, run_dir=args.run_dir, resume=args.resume) |
|
|
|
|
| def cmd_sentiment_prepare(args): |
| from environments.sentiment.dataset import prepare |
| print(json.dumps(prepare(args.out, source=args.source), indent=2)) |
|
|
|
|
| |
| def build_parser() -> argparse.ArgumentParser: |
| p = argparse.ArgumentParser(prog="minijev", description="MiniJEV: tiny calibrated decision model") |
| sub = p.add_subparsers(dest="cmd", required=True) |
|
|
| s = sub.add_parser("hardware", help="show detected device and runtime profile") |
| s.add_argument("--batch-size", type=int, default=256) |
| s.add_argument("--params", type=int, default=1_000_000) |
| s.set_defaults(fn=cmd_hardware) |
|
|
| s = sub.add_parser("train", help="supervised / distillation training") |
| s.add_argument("--config", required=True) |
| s.add_argument("--set", nargs="*", help="override config values, e.g. training.lr=5e-4") |
| s.add_argument("--run-dir") |
| s.add_argument("--resume", help="run directory to resume from (uses checkpoints/last.pt)") |
| s.add_argument("--max-steps", type=int) |
| s.add_argument("--max-minutes", type=float, help="stop cleanly (checkpoint saved) after this wall time") |
| s.add_argument("--no-games", action="store_true", help="skip game evaluation after training") |
| s.set_defaults(fn=cmd_train) |
|
|
| s = sub.add_parser("evaluate", help="calibration on a split + (snake) full games") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--data", help="dataset dir (default: from checkpoint config)") |
| s.add_argument("--split", default="test") |
| s.add_argument("--games", type=int, default=200) |
| s.add_argument("--board-size", type=int) |
| s.add_argument("--threads", type=int, default=4) |
| s.add_argument("--out") |
| s.set_defaults(fn=cmd_evaluate) |
|
|
| s = sub.add_parser("benchmark", help="latency / throughput / CPU / RAM") |
| s.add_argument("--checkpoint", required=True, help=".pt or .onnx") |
| s.add_argument("--onnx", action="store_true", help="also benchmark sibling .onnx and .int8.onnx files") |
| s.add_argument("--threads", type=int, default=1) |
| s.add_argument("--quick", action="store_true") |
| s.add_argument("--out") |
| s.set_defaults(fn=cmd_benchmark) |
|
|
| s = sub.add_parser("export", help="export to ONNX (+ INT8) and verify parity") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--format", default="onnx") |
| s.add_argument("--out") |
| s.add_argument("--no-quantize", action="store_true") |
| s.set_defaults(fn=cmd_export) |
|
|
| s = sub.add_parser("estimate", help="estimate training time and cost (measures real throughput)") |
| s.add_argument("--config", required=True) |
| s.add_argument("--set", nargs="*") |
| s.add_argument("--gpu-hourly-cost", type=float, default=0.0) |
| s.add_argument("--gpu-name") |
| s.add_argument("--throughput", type=float, help="samples/s; skip the measurement") |
| s.add_argument("--bench-seconds", type=float, default=60) |
| s.add_argument("--budget", type=float, help="warn if the estimate exceeds this amount") |
| s.set_defaults(fn=cmd_estimate) |
|
|
| s = sub.add_parser("compare", help="results table across runs") |
| s.add_argument("runs", nargs="+") |
| s.add_argument("--out", help="write markdown table here") |
| s.set_defaults(fn=cmd_compare) |
|
|
| runs = sub.add_parser("runs", help="run bookkeeping").add_subparsers(dest="runs_cmd", required=True) |
| s = runs.add_parser("annotate") |
| s.add_argument("run") |
| s.add_argument("--actual-cost", type=float) |
| s.add_argument("--gpu") |
| s.add_argument("--note") |
| s.set_defaults(fn=cmd_runs_annotate) |
|
|
| snake = sub.add_parser("snake", help="Snake environment tools").add_subparsers(dest="snake_cmd", required=True) |
| s = snake.add_parser("generate", help="generate an oracle-labelled dataset") |
| s.add_argument("--samples", type=int, default=100_000) |
| s.add_argument("--board-size", default="10", help="e.g. 10 or 10,20") |
| s.add_argument("--out") |
| s.add_argument("--seed", type=int, default=0) |
| s.add_argument("--keep-prob", type=float, default=0.1) |
| s.add_argument("--explore-eps", type=float, default=0.05) |
| s.add_argument("--workers", type=int) |
| s.add_argument("--jsonl", action="store_true", help="also write full JSONL exports") |
| s.set_defaults(fn=cmd_snake_generate) |
|
|
| s = snake.add_parser("baselines", help="play the scripted baselines") |
| s.add_argument("--games", type=int, default=100) |
| s.add_argument("--board-size", type=int, default=10) |
| s.add_argument("--seed", type=int, default=12345) |
| s.add_argument("--out") |
| s.set_defaults(fn=cmd_snake_baselines) |
|
|
| for name, human in (("demo", False), ("play", True)): |
| s = snake.add_parser(name, help="web UI: watch MiniJEV/baselines" if not human else "web UI: play yourself") |
| s.add_argument("--model", help=".onnx or .pt") |
| s.add_argument("--agent", default="minijev", help="minijev | oracle | bfs | survival | greedy | random | human") |
| s.add_argument("--board-size", type=int, default=10) |
| s.add_argument("--port", type=int, default=8765) |
| s.add_argument("--no-browser", action="store_true") |
| s.set_defaults(fn=lambda a, h=human: cmd_snake_demo(a, h)) |
|
|
| s = snake.add_parser("dagger", help="DAgger: label the model's own states with the oracle, retrain") |
| s.add_argument("--config", required=True) |
| s.add_argument("--set", nargs="*") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--iterations", type=int, default=2) |
| s.add_argument("--samples", type=int, default=100_000) |
| s.add_argument("--run-dir") |
| s.set_defaults(fn=cmd_snake_dagger) |
|
|
| s = snake.add_parser("rl", help="PPO fine-tuning from a supervised checkpoint") |
| s.add_argument("--config", required=True) |
| s.add_argument("--set", nargs="*") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--run-dir") |
| s.add_argument("--resume", action="store_true") |
| s.set_defaults(fn=cmd_snake_rl) |
|
|
| s = sub.add_parser("relabel", help="rename choice labels of a trained model (fits the ChoiceEncoder only)") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--out", required=True) |
| s.add_argument("--map", nargs="+", required=True, help='OLD="new label" ...') |
| s.add_argument("--set", nargs="*", help="config updates stored in the new checkpoint") |
| s.add_argument("--steps", type=int, default=3000) |
| s.set_defaults(fn=cmd_relabel) |
|
|
| s = sub.add_parser("export-hf", help="write a pickle-free folder (safetensors + JSON) for Hugging Face") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--out", required=True) |
| s.set_defaults(fn=cmd_export_hf) |
|
|
| s = sub.add_parser("ask", help="one decision from the command line") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--state", default="") |
| s.add_argument("--question", required=True) |
| s.add_argument("--choices", nargs="+", required=True) |
| s.add_argument("--task") |
| s.add_argument("--threads", type=int, default=1) |
| s.set_defaults(fn=cmd_ask) |
|
|
| s = sub.add_parser("bench-jsonl", help="accuracy/calibration on a JSONL decision file") |
| s.add_argument("--checkpoint", required=True) |
| s.add_argument("--data", required=True) |
| s.add_argument("--task") |
| s.add_argument("--threads", type=int, default=4) |
| s.add_argument("--out") |
| s.add_argument("-v", "--verbose", action="store_true") |
| s.set_defaults(fn=cmd_bench_jsonl) |
|
|
| sent = sub.add_parser("sentiment", help="sentiment task tools").add_subparsers(dest="sent_cmd", required=True) |
| s = sent.add_parser("prepare", help="download / build the sentiment dataset") |
| s.add_argument("--out", default="data/sentiment") |
| s.add_argument("--source", default="tweet_eval") |
| s.set_defaults(fn=cmd_sentiment_prepare) |
| return p |
|
|
|
|
| def main(argv=None): |
| args = build_parser().parse_args(argv) |
| args.fn(args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
