File size: 2,955 Bytes
1e05592 | 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 | #!/usr/bin/env python3
"""
predict_nexar_test.py
═══════════════════════════════════════════════════════════════════════════════
Generate a Kaggle-style submission CSV from a Nexar head checkpoint and the
test belief cache.
Inputs:
--head_ckpt checkpoints/Nexar/qwen3vl4b_head/best.pt
--test_cache data/belief_cache_nexar_qwen3vl4b/test.pt
Output CSV columns: id,score
Usage
─────
python -m training.Policy.predict_nexar_test \
--head_ckpt checkpoints/Nexar/qwen3vl4b_head/best.pt \
--test_cache data/belief_cache_nexar_qwen3vl4b/test.pt \
--out submissions/nexar_qwen3vl4b.csv
"""
from __future__ import annotations
import argparse
import csv
import logging
from pathlib import Path
import numpy as np
import torch
from training.Policy.train_nexar_head import NexarHead
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("Policy.predict_nexar_test")
def main():
ap = argparse.ArgumentParser("predict_nexar_test")
ap.add_argument("--head_ckpt", required=True)
ap.add_argument("--test_cache", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--batch_size", type=int, default=128)
args = ap.parse_args()
logger.info(f"loading head {args.head_ckpt}")
ck = torch.load(args.head_ckpt, map_location="cpu", weights_only=False)
meta = ck["meta"]
model = NexarHead(hidden_dim=meta["hidden_dim"],
proj_dim=meta["proj_dim"],
n_layers=meta["n_layers"],
n_heads=meta["n_heads"],
dropout=meta["dropout"])
model.load_state_dict(ck["state_dict"])
model.eval().to("cuda")
logger.info(f"loading test cache {args.test_cache}")
te = torch.load(args.test_cache, map_location="cpu", weights_only=False)
x = te["beliefs_frame"].float()
v = te["valid_frames"].bool()
ids = te["meta"]["video_ids"]
assert x.shape[0] == len(ids), f"cache/ids mismatch: {x.shape[0]} vs {len(ids)}"
probs = []
with torch.no_grad():
for i in range(0, x.size(0), args.batch_size):
xb = x[i:i + args.batch_size].to("cuda")
vb = v[i:i + args.batch_size].to("cuda")
logits = model(xb, vb).cpu().numpy()
probs.append(1 / (1 + np.exp(-logits)))
probs = np.concatenate(probs)
assert len(probs) == len(ids)
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
with open(out, "w", newline="") as f:
w = csv.writer(f)
w.writerow(["id", "score"])
for vid, p in zip(ids, probs):
w.writerow([vid, f"{float(p):.6f}"])
logger.info(f"wrote {len(ids)} rows -> {out}")
if __name__ == "__main__":
main()
|