| """Score checkpoints on one common, larger validation set (same windows for every run). |
| Usage: source env.sh && $TA_PY scripts/rescore.py run_dir [run_dir ...] [--batches 12]""" |
| import argparse |
| import json |
| import os |
|
|
| import numpy as np |
| import torch |
|
|
| from tiny_agent.checkpoint import load_model |
| from tiny_agent.data import MixtureLoader |
| from tiny_agent.model import make_block_mask |
|
|
|
|
| @torch.no_grad() |
| def score(model, batches): |
| model.eval() |
| out = {} |
| for name, bs in batches.items(): |
| ls = [] |
| for inp, tgt, doc in bs: |
| inp, tgt, doc = inp.to("xpu"), tgt.to("xpu"), doc.to("xpu") |
| with torch.autocast("xpu", dtype=torch.bfloat16): |
| ls.append(model(inp, doc, make_block_mask(doc, model.cfg.swa_window), tgt).item()) |
| out[name] = float(np.mean(ls)) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("runs", nargs="+") |
| ap.add_argument("--batches", type=int, default=12) |
| a = ap.parse_args() |
| batches = MixtureLoader("val", "stable", 2048, 8, stream=False).fixed_batches(a.batches, seed=4321) |
| res = {} |
| for r in a.runs: |
| st = torch.load(os.path.join(r, "ckpt.pt"), map_location="cpu", weights_only=False) |
| model = load_model(os.path.join(r, "ckpt.pt")) |
| v = score(model, batches) |
| res[r] = {"tokens_M": round(st["tokens"] / 1e6), "step": st["step"], "mean_val": round(float(np.mean(list(v.values()))), 4), |
| **{k: round(x, 4) for k, x in sorted(v.items())}} |
| print(json.dumps({r: res[r]}), flush=True) |
| del model |
| torch.xpu.empty_cache() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
