| """Portable one-GPU JSONL inference for frozen dynamic-option checkpoints.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
| import time |
| import torch |
| from decision_model import DecisionModel, encode, collate |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--model", required=True) |
| p.add_argument("--base", action="store_true") |
| p.add_argument("--revision", default="local-checkpoint") |
| p.add_argument("--input", required=True) |
| p.add_argument("--output", required=True) |
| p.add_argument("--batch-size", type=int, default=8) |
| p.add_argument("--max-length", type=int, default=4096) |
| p.add_argument("--temperature", type=float, default=1.0) |
| a = p.parse_args() |
| assert a.temperature > 0 |
| torch.cuda.set_device(0) |
| model, tokenizer = DecisionModel.from_base(a.model, a.revision) if a.base else DecisionModel.from_checkpoint(a.model) |
| model = model.cuda().eval() |
| rows = [json.loads(line) for line in Path(a.input).read_text().splitlines() if line.strip()] |
| encoded = [encode(row, tokenizer, a.max_length) for row in rows] |
| assert len({row["id"] for row in rows}) == len(rows) |
| output = Path(a.output); output.parent.mkdir(parents=True, exist_ok=True) |
| if output.exists(): raise RuntimeError("Refusing to overwrite predictions") |
| with torch.inference_mode(), output.open("w") as f: |
| for start in range(0, len(encoded), a.batch_size): |
| examples = encoded[start:start + a.batch_size] |
| batch = {key: value.cuda() if torch.is_tensor(value) else value for key, value in collate(examples, tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id).items()} |
| torch.cuda.synchronize(); tick = time.perf_counter() |
| with torch.autocast("cuda", dtype=torch.bfloat16): logits = model(**batch) |
| torch.cuda.synchronize(); elapsed = time.perf_counter() - tick |
| probabilities = (logits / a.temperature).softmax(-1).cpu().tolist(); scores = logits.cpu().tolist() |
| for row, example, prob, score in zip(rows[start:start + a.batch_size], examples, probabilities, scores): |
| k = example["nopts"]; pred = max(range(k), key=lambda i: prob[i]) |
| rec = {"id": row["id"], "family": row.get("family"), "label": row.get("label"), "prediction": pred, "prediction_key": row["options"][pred]["key"], "probabilities": prob[:k], "logits": score[:k], "temperature": a.temperature, "input_tokens": len(example["ids"]), "prompt_sha256": example["prompt_sha256"], "batch_elapsed_seconds": elapsed, "batch_size": len(examples)} |
| if "score_values" in row: rec["expected_score"] = sum(value * probability for value, probability in zip(row["score_values"], prob[:k])) |
| f.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| f.flush(); print(json.dumps({"completed": min(start + a.batch_size, len(rows)), "total": len(rows)}), flush=True) |
| Path(str(output) + ".metadata.json").write_text(json.dumps({"model": a.model, "revision": a.revision, "temperature": a.temperature, "input_sha256": hashlib.sha256(Path(a.input).read_bytes()).hexdigest(), "predictions_sha256": hashlib.sha256(output.read_bytes()).hexdigest(), "rows": len(rows), "model_metadata": model.metadata}, indent=2)) |
|
|
