code/minijev/cli.py

"""`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 [])


# ------------------------------------------------------------------ generic
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)  # T already applied
        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))  # logits are already calibrated
    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))


# -------------------------------------------------------------------- snake
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))


# --------------------------------------------------------------------- main
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()