| """Predict with a trained run. Run as ``python -m training.predict`` (pod, or a rented GPU for backfills). |
| |
| Input: a bundle split (``--split val|test``) or any JSONL of rows with ``id`` and ``segments`` (``--input``: the |
| robustness set, or ``naver-psy predict --export-inputs`` for a backfill, classifier-plan §6.4). Output rows: |
| ``{"id", "split", "n_tokens", "truncated", "logits"}`` with raw logits per field in taxonomy option order (ids and |
| numbers only); ``training.calibrate`` and the predict worker turn them into calibrated, gated probabilities. |
| The summary line reports posts/s and tokens/s and, with ``--repeat N``, the largest calibrated probability |
| difference between repeats (the determinism check, §5.5 item 6). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import time |
| from pathlib import Path |
|
|
| from .calibrate import max_prob_diff, outputs |
| from .data import inference_batches, load_jsonl, load_split, write_jsonl |
| from .formats import encode |
| from .taxonomy import Taxonomy |
|
|
|
|
| def load_run(run: Path, which: str = "best", model_dir: str | None = None, device="cuda"): |
| """(model, tokenizer, taxonomy, config) for a run directory.""" |
| from .model import build, load_tokenizer |
|
|
| cfg = json.loads((run / "config.json").read_text()) |
| if model_dir: |
| cfg["model"] = model_dir |
| tax = Taxonomy.load(run / "taxonomy.json") |
| ckpt = run / which |
| model = build(cfg, tax.sizes(), device, train=False, adapter_dir=ckpt / "adapter", |
| heads_path=ckpt / "heads.safetensors") |
| return model, load_tokenizer(cfg["model"]), tax, cfg |
|
|
|
|
| def main(argv=None) -> int: |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| p.add_argument("--run", required=True, help="Run directory written by training.train") |
| p.add_argument("--which", default="best", choices=["best", "last"]) |
| source = p.add_mutually_exclusive_group(required=True) |
| source.add_argument("--split", choices=["train", "val", "test"]) |
| source.add_argument("--input", help="JSONL of rows with id and segments") |
| p.add_argument("--data", help="Bundle directory (default: the run config's)") |
| p.add_argument("--model", help="Base checkpoint directory (default: the run config's)") |
| p.add_argument("--out", required=True) |
| p.add_argument("--token-budget", type=int, default=65536, help="Padded tokens per batch") |
| p.add_argument("--limit", type=int, default=0) |
| p.add_argument("--repeat", type=int, default=1) |
| p.add_argument("--device", default="cuda") |
| args = p.parse_args(argv) |
| import torch |
|
|
| from .model import bos_id, predict, tokenizer_fn |
|
|
| run = Path(args.run) |
| device = torch.device(args.device) |
| model, tok, tax, cfg = load_run(run, args.which, args.model, device) |
| data = args.data or cfg["data"] |
| rows = load_split(data, args.split) if args.split else load_jsonl(args.input) |
| if args.limit: |
| rows = rows[:args.limit] |
| t = time.time() |
| encoded = [encode(r["segments"], tokenizer_fn(tok), cfg["max_len"], bos_id(tok)) for r in rows] |
| ids = [e[0] for e in encoded] |
| encode_seconds = time.time() - t |
| batches = inference_batches([len(x) for x in ids], args.token_budget) |
| runs = [] |
| for _ in range(max(1, args.repeat)): |
| t = time.time() |
| logits = predict(model, ids, tok.pad_token_id, batches, device, tax.names) |
| if device.type == "cuda": |
| torch.cuda.synchronize() |
| runs.append((logits, time.time() - t)) |
| logits, seconds = runs[0] |
| out_rows = [{"id": r["id"], "split": r.get("split"), "n_tokens": encoded[i][1]["n_tokens"], |
| "truncated": encoded[i][1]["cut"], "logits": logits[i]} for i, r in enumerate(rows)] |
| write_jsonl(args.out, out_rows) |
| tokens = sum(len(x) for x in ids) |
| summary = dict(rows=len(rows), seconds=round(seconds, 2), encode_seconds=round(encode_seconds, 2), |
| posts_per_second=round(len(rows) / max(seconds, 1e-9), 1), |
| tokens_per_second=round(tokens / max(seconds, 1e-9)), tokens=tokens, |
| truncated_rows=sum(1 for e in encoded if e[1]["cut"]), which=args.which) |
| if len(runs) > 1: |
| temps_path = run / "temperatures.json" |
| temps = json.loads(temps_path.read_text()) if temps_path.exists() else None |
| diff, same = 0.0, True |
| for other, _ in runs[1:]: |
| for i in logits: |
| a, b = outputs(logits[i], tax, temps), outputs(other[i], tax, temps) |
| diff = max(diff, max_prob_diff(a["probs"], b["probs"])) |
| same &= a["labels"] == b["labels"] |
| summary.update(max_prob_diff=diff, argmax_identical=same, deterministic=diff <= 1e-3) |
| if device.type == "cuda": |
| summary["max_memory_gb"] = round(torch.cuda.max_memory_allocated() / 2 ** 30, 1) |
| print(json.dumps(summary)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
