| """N1: realized acceptance of saved drafts (native_parity *_drafts.pt) against the target's own tokens in the |
| captures, paired with the engine's drafts of the same blocks, per window (A / B) and category. CPU only. Aggregates |
| only. Several drafts files: each arm is paired with the engine and with the first file (the reference arm). |
| Usage (x86host): python -m offload.dspark.native.tools.native_accept OUT.json DRAFTS.pt [DRAFTS.pt ...] |
| """ |
| from __future__ import annotations |
|
|
| import collections |
| import json |
| import math |
| import sys |
|
|
| import torch |
|
|
| from ...worker.tools.parity_capture import category |
| from .native_parity import CAPS, accept_len, load_rounds, realized_tokens |
|
|
|
|
| def paired(a: list[int], b: list[int]) -> dict: |
| n = len(a) |
| if not n: |
| return {"n": 0} |
| d = [x - y for x, y in zip(a, b)] |
| m = sum(d) / n |
| sd = math.sqrt(sum((x - m) ** 2 for x in d) / max(n - 1, 1)) |
| return {"n": n, "mean_a": round(sum(a) / n, 4), "mean_b": round(sum(b) / n, 4), "mean_diff": round(m, 4), |
| "se_diff": round(sd / math.sqrt(n), 4), "a_better": sum(x > 0 for x in d), "b_better": sum(x < 0 for x in d)} |
|
|
|
|
| def main() -> None: |
| out_path, files = sys.argv[1], sys.argv[2:] |
| rounds = load_rounds(CAPS) |
| real = realized_tokens(rounds) |
| by_round = {int(r["round"]): r for r in rounds} |
| runs = {f.rsplit("/", 1)[-1].replace("_drafts.pt", ""): torch.load(f, weights_only=False) for f in files} |
| names = list(runs) |
| L: dict = collections.defaultdict(list) |
| keys = [] |
| for idx, recs in enumerate(zip(*runs.values())): |
| rnd = recs[0]["round"] |
| if any(r["round"] != rnd for r in recs): |
| raise SystemExit("drafts files replayed different rounds") |
| r = by_round[rnd] |
| req_ids = list(r["req_ids"]) |
| nq = int(r["num_query_per_req"]) |
| qpos = r["query_positions"].view(len(req_ids), nq).long() |
| eng = r["draft_tokens"].long() |
| nsam = r["num_sampled"].tolist() |
| for b, req in enumerate(req_ids): |
| if category(req, nsam[b]) != "live": |
| continue |
| first = int(qpos[b, 0]) + 1 |
| le, oke = accept_len(eng[b].tolist(), first, real[req]) |
| for w in "AB": |
| vals = [accept_len(rec[w][b], first, real[req]) for rec in recs] |
| if not oke or not all(ok for _, ok in vals): |
| continue |
| L[("engine", w)].append(le) |
| for name, (la, _) in zip(names, vals): |
| L[(name, w)].append(la) |
| rep = {"captures": CAPS, "arms": names, "category": "live"} |
| for w in "AB": |
| rep[w] = {"blocks": len(L[("engine", w)]), |
| "engine_mean_accept": round(sum(L[("engine", w)]) / max(1, len(L[("engine", w)])), 4)} |
| for name in names: |
| rep[w][name] = {"vs_engine": paired(L[(name, w)], L[("engine", w)]), |
| "vs_first_arm": paired(L[(name, w)], L[(names[0], w)]) if name != names[0] else None} |
| with open(out_path, "w") as f: |
| json.dump(rep, f, indent=1) |
| print(json.dumps(rep, indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
