File size: 4,690 Bytes
60b165e | 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 | #!/usr/bin/env python3
"""
Tool-selection benchmark for AUTOREGRESSIVE models via mlx_lm — the apples-to-apples
counterpart to diffusion_eval.py. Same clean test split, same parse_tools, same
metrics, so any AR model's number is directly comparable to DiffusionGemma's.
It reconstructs the raw system/user text from the DiffusionGemma-rendered prompt,
then re-renders it with the target model's OWN chat template (fair zero-shot framing).
Usage:
python3 ar_eval.py \
--model mlx-community/Qwen3.6-35B-A3B-4bit \
--test ./data/test.jsonl --out ./eval_qwen.json \
--label Qwen3.6-35B [--max-samples 0]
"""
import argparse, json, re, time
from pathlib import Path
import mlx.core as mx
from mlx_lm import load, generate
from mlx_lm.sample_utils import make_sampler
def parse_tools(text):
tools = []
for raw in text.split("\n"):
line = raw.strip()
for stop in ("<turn|>", "<eos>", "<|turn>", "</s>", "<|im_end|>", "<end_of_turn>"):
if stop in line:
line = line.split(stop)[0].strip()
# accept "- Tool", "1. Tool", "* Tool", or bare "Tool" on its own line
m = re.match(r"^[\-\*\d.\)\s]*([A-Za-z][A-Za-z0-9_\-.]+)\s*$", line)
if line.startswith(("- ", "* ")) or re.match(r"^\d+[.)]", line):
if m:
tools.append(m.group(1))
seen, out = set(), []
for t in tools:
if t not in seen:
seen.add(t); out.append(t)
return out
def metrics(pred, expected):
ps, es = set(pred), set(expected)
inter, union = ps & es, ps | es
return dict(
jaccard=len(inter)/len(union) if union else 1.0,
exact=1.0 if ps == es else 0.0,
precision=len(inter)/len(ps) if ps else 0.0,
recall=len(inter)/len(es) if es else 0.0,
top1=1.0 if pred and pred[0] in es else 0.0,
)
def reconstruct(prompt_text):
"""Pull raw system + user back out of the DiffusionGemma-rendered prompt."""
sys_txt = prompt_text.split("<|turn>system\n", 1)[1].split("<turn|>", 1)[0].strip() \
if "<|turn>system\n" in prompt_text else ""
usr_txt = prompt_text.split("<|turn>user\n", 1)[1].split("<turn|>", 1)[0].strip()
return sys_txt, usr_txt
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--test", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--label", required=True)
ap.add_argument("--max-samples", type=int, default=0)
ap.add_argument("--max-tokens", type=int, default=96)
args = ap.parse_args()
print(f"[load] {args.model}", flush=True)
model, tokenizer = load(args.model)
sampler = make_sampler(temp=0.0) # greedy, deterministic
samples = []
with open(args.test) as f:
for line in f:
o = json.loads(line)
sys_txt, usr_txt = reconstruct(o["prompt"])
expected = parse_tools(o["response"].split("<turn|>")[0])
samples.append((sys_txt, usr_txt, expected))
if args.max_samples:
samples = samples[: args.max_samples]
print(f"[data] {len(samples)} samples", flush=True)
results, per_sample = [], []
t0 = time.time()
for i, (sys_txt, usr_txt, expected) in enumerate(samples):
msgs = ([{"role": "system", "content": sys_txt}] if sys_txt else []) + \
[{"role": "user", "content": usr_txt}]
try:
prompt = tokenizer.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
except Exception:
prompt = (sys_txt + "\n\n" if sys_txt else "") + usr_txt
text = generate(model, tokenizer, prompt=prompt, max_tokens=args.max_tokens,
sampler=sampler, verbose=False)
pred = parse_tools(text)
m = metrics(pred, expected)
results.append(m)
per_sample.append({"i": i, "pred": pred, "expected": expected, "raw": text[:300], **m})
mx.clear_cache() # bound MLX buffer cache over a long eval (see diffusion_eval.py)
if (i + 1) % 20 == 0:
print(f" {i+1}/{len(samples)} jac so far="
f"{sum(r['jaccard'] for r in results)/len(results):.4f} "
f"({(time.time()-t0)/(i+1):.1f}s/sample)", flush=True)
n = len(results)
agg = {k: sum(r[k] for r in results)/n for k in results[0]}
agg.update(n=n, model=args.model, label=args.label, seconds_total=round(time.time()-t0, 1))
print(json.dumps(agg, indent=2), flush=True)
Path(args.out).write_text(json.dumps({"aggregate": agg, "samples": per_sample}, indent=2))
print(f"[done] {args.out}", flush=True)
if __name__ == "__main__":
main()
|