"""Benchmark a model on the UnchaosLM triage test set (fully synthetic). Scores two capabilities: 1. Extraction — strict-JSON triage records from raw work signals. Metrics: valid-JSON rate, accuracy on the 7 deterministic fields (is_actionable, action_direction, severity, tier, owner, deadline_suggestion, escalation_flag), escalation recall. 2. Agent — native Qwen3 tool calling. Metrics: correct tool name, exact-argument match, on rows whose gold ends in a tool call. Free-text fields (severity_rationale, why_it_matters, next_step) are not auto-scored. Greedy decoding, thinking disabled. Usage: python eval.py [test.jsonl] [results.json] """ import json import re import sys from mlx_lm import load, generate CHECKED_FIELDS = [ "is_actionable", "action_direction", "severity", "tier", "owner", "deadline_suggestion", "escalation_flag", ] MAX_TOKENS = 512 THINK_RE = re.compile(r".*?", re.DOTALL) TOOL_CALL_RE = re.compile(r"\s*(\{.*?\})\s*", re.DOTALL) # Qwen3.5-style XML tool calls: value... FUNC_XML_RE = re.compile(r"(.*?)(?:|$)", re.DOTALL) PARAM_XML_RE = re.compile(r"\s*(.*?)\s*", re.DOTALL) def parse_json_obj(text): text = THINK_RE.sub("", text) start = text.find("{") if start == -1: return None try: # first complete JSON object; ignores trailing text/repetition obj, _ = json.JSONDecoder().raw_decode(text[start:]) return obj if isinstance(obj, dict) else None except json.JSONDecodeError: return None def parse_tool_calls(text): text = THINK_RE.sub("", text) calls = [] for m in TOOL_CALL_RE.finditer(text): try: calls.append(json.loads(m.group(1))) except json.JSONDecodeError: pass if not calls: # fall back to Qwen3.5's XML function format for name, body in FUNC_XML_RE.findall(text): args = {k: v for k, v in PARAM_XML_RE.findall(body)} calls.append({"name": name, "arguments": args}) return calls def norm(v): return v.strip().lower() if isinstance(v, str) else v def coerce(v): if isinstance(v, str) and re.fullmatch(r"-?\d+", v.strip()): return int(v.strip()) return v.strip() if isinstance(v, str) else v def norm_args(a): return {k: coerce(v) for k, v in (a or {}).items()} def main(): # --skip-agent: extraction only, for models whose tool-call output format # isn't Qwen's XML (scoring them on it would be unfair) skip_agent = "--skip-agent" in sys.argv args = [a for a in sys.argv[1:] if a != "--skip-agent"] model_path = args[0] test_path = args[1] if len(args) > 1 else "triage_ds_v4/test.jsonl" out_path = args[2] if len(args) > 2 else "eval_results.json" rows = [json.loads(l) for l in open(test_path)] model, tokenizer = load(model_path) ex = dict(n=0, valid_json=0, fields_total=0, fields_correct=0, esc_gold=0, esc_caught=0, esc_false_alarms=0, per_field={f: [0, 0] for f in CHECKED_FIELDS}) ag = dict(n=0, tool_correct=0, args_correct=0, skipped_answer_rows=0) raw_log = [] for i, row in enumerate(rows): msgs, tools = row["messages"], row.get("tools") gold = msgs[-1] if tools and (skip_agent or not gold.get("tool_calls")): ag["skipped_answer_rows"] += 1 # grounded-answer rows: not auto-scorable continue prompt = tokenizer.apply_chat_template( msgs[:-1], tools=tools, add_generation_prompt=True, tokenize=False, enable_thinking=False) out = generate(model, tokenizer, prompt=prompt, max_tokens=MAX_TOKENS, verbose=False) raw_log.append({"idx": i, "kind": "agent" if tools else "extraction", "gold": gold.get("tool_calls") or gold.get("content"), "output": out}) if tools: ag["n"] += 1 gold_calls = [(c["function"]["name"], c["function"]["arguments"]) for c in gold["tool_calls"]] pred_calls = [(c.get("name"), c.get("arguments")) for c in parse_tool_calls(out)] if [n for n, _ in pred_calls] == [n for n, _ in gold_calls]: ag["tool_correct"] += 1 if [norm_args(a) for _, a in pred_calls] == [norm_args(a) for _, a in gold_calls]: ag["args_correct"] += 1 else: ex["n"] += 1 gold_obj = json.loads(gold["content"]) pred = parse_json_obj(out) if pred is not None: ex["valid_json"] += 1 for f in CHECKED_FIELDS: ex["fields_total"] += 1 ex["per_field"][f][1] += 1 if pred is not None and norm(pred.get(f)) == norm(gold_obj.get(f)): ex["fields_correct"] += 1 ex["per_field"][f][0] += 1 if gold_obj.get("escalation_flag") is True: ex["esc_gold"] += 1 if pred is not None and pred.get("escalation_flag") is True: ex["esc_caught"] += 1 elif pred is not None and pred.get("escalation_flag") is True: ex["esc_false_alarms"] += 1 done = i + 1 if done % 10 == 0 or done == len(rows): print(f"[{done}/{len(rows)}]", flush=True) results = {"model": model_path, "extraction": ex, "agent": ag} json.dump(results, open(out_path, "w"), indent=2) with open(out_path.replace(".json", "_raw.jsonl"), "w") as f: for r in raw_log: f.write(json.dumps(r) + "\n") print(f"\n=== {model_path} ===") print(f"Extraction ({ex['n']} signals):") print(f" valid JSON: {ex['valid_json']}/{ex['n']}") print(f" field accuracy: {ex['fields_correct']}/{ex['fields_total']} " f"({100 * ex['fields_correct'] / ex['fields_total']:.1f}%)") for f, (c, t) in ex["per_field"].items(): print(f" {f:20s} {c}/{t}") print(f" escalation recall: {ex['esc_caught']}/{ex['esc_gold']} " f"false alarms: {ex['esc_false_alarms']}/{ex['n'] - ex['esc_gold']}") print(f"Agent ({ag['n']} tool-call rows, {ag['skipped_answer_rows']} answer rows skipped):") print(f" correct tool: {ag['tool_correct']}/{ag['n']} exact args: {ag['args_correct']}/{ag['n']}") if __name__ == "__main__": main()