eval/eval.py

"""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 <model_path_or_hf_id> [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"<think>.*?</think>", re.DOTALL)
TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)
# Qwen3.5-style XML tool calls: <function=name><parameter=key>value</parameter>...</function>
FUNC_XML_RE = re.compile(r"<function=([\w.-]+)>(.*?)(?:</function>|$)", re.DOTALL)
PARAM_XML_RE = re.compile(r"<parameter=([\w.-]+)>\s*(.*?)\s*</parameter>", 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 <tool_call> 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()