File size: 4,569 Bytes
b2ebc95 | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | #!/usr/bin/env python3
"""Offline smoke test: does the served checkpoint drive the pi loop correctly?
No sandbox needed. Replays prefixes of held-out converted trajectories against the
served model and reports, over N samples:
* fraction of turns that emit a syntactically valid tool call
* fraction that use a real pi tool name (read/bash/edit/write)
* completion length distribution and finish_reason (i.e. does it stop?)
* fraction that hallucinate a tool response / continue past its own turn
Usage: check_policy.py --base-url http://localhost:8000/v1 --model policy -n 40
"""
import argparse
import collections
import glob
import json
import random
import re
import sys
import urllib.request
PI_TOOLS = {"read", "bash", "edit", "write"}
LEAK = re.compile(r"<\|im_start\|>|<\|im_end\|>|<tool_response>")
def post(url, payload, timeout=300):
req = urllib.request.Request(
url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}
)
with urllib.request.urlopen(req, timeout=timeout) as r:
return json.load(r)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--base-url", default="http://localhost:8000/v1")
ap.add_argument("--model", default="policy")
ap.add_argument("--data", default="/mnt/pvc/users/simon/agentptb/runs/a-opus-max/workspace/data/sft_v1")
ap.add_argument("-n", type=int, default=40)
ap.add_argument("--temperature", type=float, default=0.7)
ap.add_argument("--max-tokens", type=int, default=2048)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
random.seed(args.seed)
import pyarrow.parquet as pq
f = sorted(glob.glob(args.data + "/*.parquet"))[-1]
rows = pq.ParquetFile(f).read_row_group(0).to_pylist()
random.shuffle(rows)
stats = collections.Counter()
lens = []
for row in rows[: args.n]:
msgs = row["messages"]
tools = json.loads(row["tool_defs"])
# cut at a random assistant turn: the model must produce that turn
idxs = [i for i, m in enumerate(msgs) if m["role"] == "assistant"]
if len(idxs) < 2:
continue
cut = random.choice(idxs[:-1])
prefix = []
for m in msgs[:cut]:
mm = {k: v for k, v in m.items() if k != "reasoning_content" and v not in (None, "")}
if m["role"] == "assistant" and m.get("reasoning_content") and not m.get("content"):
mm["content"] = m["reasoning_content"]
prefix.append(mm)
if not prefix or prefix[-1]["role"] == "assistant":
continue
try:
out = post(
args.base_url + "/chat/completions",
{
"model": args.model,
"messages": prefix,
"tools": tools,
"temperature": args.temperature,
"max_tokens": args.max_tokens,
},
)
except Exception as e:
stats["http_error"] += 1
print("ERR", repr(e)[:200], file=sys.stderr)
continue
ch = out["choices"][0]
msg = ch["message"]
stats["n"] += 1
stats["finish_" + str(ch.get("finish_reason"))] += 1
lens.append(out["usage"]["completion_tokens"])
tcs = msg.get("tool_calls") or []
if tcs:
stats["has_tool_call"] += 1
for tc in tcs:
name = (tc.get("function") or {}).get("name")
stats["tool_" + str(name)] += 1
if name in PI_TOOLS:
stats["valid_tool_name"] += 1
try:
json.loads((tc.get("function") or {}).get("arguments") or "{}")
stats["valid_args"] += 1
except Exception:
stats["bad_args"] += 1
else:
stats["text_only"] += 1
text = msg.get("content") or ""
if LEAK.search(text):
stats["leaked_special_tokens"] += 1
lens.sort()
n = max(1, stats["n"])
print(json.dumps({k: v for k, v in sorted(stats.items())}, indent=1))
if lens:
print(
f"completion_tokens: mean={sum(lens)//len(lens)} median={lens[len(lens)//2]} "
f"p90={lens[int(0.9*len(lens))]} max={lens[-1]}"
)
print(f"tool-call rate {stats['has_tool_call']/n:.2f} valid-name rate {stats['valid_tool_name']/max(1,stats['has_tool_call']):.2f} leak rate {stats['leaked_special_tokens']/n:.2f}")
if __name__ == "__main__":
main()
|