File size: 4,913 Bytes
210f9af | 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 | """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())
|