elizaos--eliza/packages/training/scripts/benchmark/native_tool_call_bench.py

"""Benchmark native Eliza trajectory rows for tool-call structure.

The input is `eliza_native_v1` JSONL (legacy flat `ElizaRecord` rows and
chatml `{messages,tools}` rows are coerced to that shape automatically — see
_to_native). For each row, the benchmark renders the
request side with the tokenizer chat template, generates from a base or tuned
Gemma / Eliza-1 checkpoint, and compares the decoded output to the native
response side.

Primary score:
  - tool_call_structure: expected native tool names and argument keys appear
  - json_structure: non-tool JSON tasks keep the expected decision/action shape
"""

from __future__ import annotations

import argparse
import json
import logging
import re
import sys
import time
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(ROOT / "scripts"))

from format_for_training import format_record  # noqa: E402
from lib.attn import select_attn_impl  # noqa: E402

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger("native-bench")

JSON_FENCE_RE = re.compile(r"^```(?:json)?\s*([\s\S]*?)\s*```$", re.I)


@dataclass
class BucketResult:
    name: str
    n: int = 0
    structure_ok: int = 0
    content_ok: int = 0
    parse_errors: int = 0
    field_match: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    field_total: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    failures: list[dict[str, Any]] = field(default_factory=list)
    # Throughput accounting — wallclock seconds spent inside model.generate()
    # for this bucket, plus the prompt / generated token counts seen.
    gen_seconds: float = 0.0
    prompt_tokens: int = 0
    gen_tokens: int = 0

    def record(
        self,
        *,
        ok_structure: bool,
        ok_content: bool,
        parse_error: bool,
        fields: dict[str, bool],
        failed: dict[str, Any] | None,
        gen_dt: float = 0.0,
        n_prompt_tokens: int = 0,
        n_gen_tokens: int = 0,
    ) -> None:
        self.n += 1
        self.structure_ok += int(ok_structure)
        self.content_ok += int(ok_content)
        self.parse_errors += int(parse_error)
        self.gen_seconds += gen_dt
        self.prompt_tokens += n_prompt_tokens
        self.gen_tokens += n_gen_tokens
        for key, ok in fields.items():
            self.field_total[key] += 1
            self.field_match[key] += int(ok)
        if failed and len(self.failures) < 8:
            self.failures.append(failed)

    def to_dict(self) -> dict[str, Any]:
        def pct(num: int, denom: int) -> float:
            return round(100.0 * num / denom, 2) if denom else 0.0
        def rate(num: int, denom: float) -> float:
            return round(num / denom, 2) if denom > 0 else 0.0

        return {
            "bucket": self.name,
            "n": self.n,
            "structure_ok": self.structure_ok,
            "structure_pct": pct(self.structure_ok, self.n),
            "content_ok": self.content_ok,
            "content_pct": pct(self.content_ok, self.n),
            "parse_errors": self.parse_errors,
            # prompt_tps / gen_tps share the same denominator (the wallclock
            # spent inside model.generate(): prefill + decode); model.generate
            # does not separate the two phases, so these are throughput proxies.
            "prompt_tps": rate(self.prompt_tokens, self.gen_seconds),
            "gen_tps": rate(self.gen_tokens, self.gen_seconds),
            "gen_seconds": round(self.gen_seconds, 2),
            "field_match_pct": {
                key: pct(self.field_match[key], total)
                for key, total in self.field_total.items()
            },
            "failures": self.failures,
        }


def _clean_json_text(text: str) -> str:
    stripped = text.strip()
    match = JSON_FENCE_RE.match(stripped)
    if match:
        return match.group(1).strip()
    return stripped


def _parse_json_object(text: str) -> dict[str, Any] | None:
    cleaned = _clean_json_text(text)
    start = cleaned.find("{")
    if start < 0:
        return None
    # Walk forward to find the balanced closing brace. Using rfind would
    # corrupt the slice when the model appends trailing tool-call markup after
    # the eliza JSON object.
    depth = 0
    in_str = False
    escape = False
    end = -1
    for i, ch in enumerate(cleaned[start:], start):
        if escape:
            escape = False
            continue
        if ch == "\\" and in_str:
            escape = True
            continue
        if ch == '"':
            in_str = not in_str
            continue
        if in_str:
            continue
        if ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1
            if depth == 0:
                end = i
                break
    if end < 0:
        return None
    try:
        value = json.loads(cleaned[start:end + 1])
    except json.JSONDecodeError:
        return None
    return value if isinstance(value, dict) else None


def _call_name(call: dict[str, Any]) -> str:
    function = call.get("function") if isinstance(call.get("function"), dict) else {}
    value = call.get("toolName") or call.get("name") or function.get("name")
    return value if isinstance(value, str) else ""


def _call_args(call: dict[str, Any]) -> dict[str, Any]:
    function = call.get("function") if isinstance(call.get("function"), dict) else {}
    args = (
        call.get("input")
        if "input" in call
        else call.get("args")
        if "args" in call
        else call.get("arguments")
        if "arguments" in call
        else function.get("arguments")
    )
    if isinstance(args, str):
        try:
            parsed = json.loads(args)
        except json.JSONDecodeError:
            return {}
        return parsed if isinstance(parsed, dict) else {}
    return args if isinstance(args, dict) else {}


def normalize_tool_calls(value: Any) -> list[dict[str, Any]]:
    if not isinstance(value, list):
        return []
    calls: list[dict[str, Any]] = []
    for raw in value:
        if isinstance(raw, dict) and _call_name(raw):
            calls.append({"name": _call_name(raw), "arguments": _call_args(raw)})
    return calls


# Gemma-4 native tool-call surface syntax emitted by a model trained on this
# chat template: `<|tool_call>call:NAME{key:<|"|>value<|"|>,...}<tool_call|>`.
# String values are wrapped in `<|"|>` markers; numbers/booleans/nested
# objects/arrays are bare. This is what the fine-tuned checkpoint produces —
# NOT JSON — so scoring native tool-calling requires parsing it directly.
_NATIVE_STR_MARK = '<|"|>'
_NATIVE_TC_RE = re.compile(
    r"<\|tool_call>\s*call:\s*([A-Za-z0-9_.\-]+)\s*\{(.*?)\}\s*<tool_call\|>", re.S
)
_NATIVE_TC_OPEN_RE = re.compile(
    r"<\|tool_call>\s*call:\s*([A-Za-z0-9_.\-]+)\s*\{(.*)$", re.S
)
_NATIVE_KEY_RE = re.compile(r"[A-Za-z0-9_.\-]+")


def _native_arg_keys(body: str) -> list[str]:
    """Extract the top-level argument keys from a native tool-call body.
    Skips `<|"|>`-quoted strings and descends past nested `{}` / `[]` so only
    depth-0 `key:` names are collected (values are not scored)."""
    keys: list[str] = []
    i, n, depth, buf, expect_key = 0, len(body), 0, "", True
    while i < n:
        if body.startswith(_NATIVE_STR_MARK, i):
            close = body.find(_NATIVE_STR_MARK, i + len(_NATIVE_STR_MARK))
            if close < 0:
                break
            i = close + len(_NATIVE_STR_MARK)
            expect_key = False
            buf = ""
            continue
        ch = body[i]
        if ch in "{[":
            depth += 1
            expect_key = False
            buf = ""
        elif ch in "}]":
            depth -= 1
        elif depth == 0 and ch == ":" and expect_key:
            key = buf.strip()
            if _NATIVE_KEY_RE.fullmatch(key):
                keys.append(key)
            expect_key = False
            buf = ""
        elif depth == 0 and ch == ",":
            expect_key = True
            buf = ""
        elif depth == 0:
            buf += ch
        i += 1
    return keys


def extract_native_tool_calls(text: str) -> list[dict[str, Any]]:
    calls: list[dict[str, Any]] = []
    for match in _NATIVE_TC_RE.finditer(text):
        keys = _native_arg_keys(match.group(2))
        calls.append({"name": match.group(1), "arguments": {k: "" for k in keys}})
    if not calls:
        # Tolerate a truncated final block with no closing `<tool_call|>`.
        match = _NATIVE_TC_OPEN_RE.search(text)
        if match:
            keys = _native_arg_keys(match.group(2))
            calls.append({"name": match.group(1), "arguments": {k: "" for k in keys}})
    return calls


def extract_tool_calls_from_text(text: str) -> list[dict[str, Any]]:
    native = extract_native_tool_calls(text)
    if native:
        return native
    parsed = _parse_json_object(text)
    if not parsed:
        return []
    for key in ("toolCalls", "tool_calls"):
        if isinstance(parsed.get(key), list):
            return normalize_tool_calls(parsed[key])
    if _call_name(parsed):
        return normalize_tool_calls([parsed])
    return []


def response_text(record: dict[str, Any]) -> str:
    response = record.get("response") if isinstance(record.get("response"), dict) else {}
    text = response.get("text")
    return text if isinstance(text, str) else ""


def expected_tool_calls(record: dict[str, Any]) -> list[dict[str, Any]]:
    response = record.get("response") if isinstance(record.get("response"), dict) else {}
    return normalize_tool_calls(response.get("toolCalls"))


def classify(record: dict[str, Any]) -> str:
    metadata = record.get("metadata") if isinstance(record.get("metadata"), dict) else {}
    task_type = metadata.get("task_type") or record.get("purpose") or "response"
    normalized = str(task_type).replace("-", "_").lower()
    if expected_tool_calls(record):
        return "tool_call"
    if normalized in {"should_respond", "context_routing"}:
        return "routing_json"
    if normalized in {"action_planner", "planner"}:
        return "planner_json"
    return "response"


def _to_native(record: dict[str, Any]) -> dict[str, Any] | None:
    """Coerce any supported corpus row into the `eliza_native_v1` shape this
    benchmark scores against. `eliza_native_v1` rows pass through. Legacy flat
    `ElizaRecord` rows and chatml `{messages,tools}` rows are run through
    `format_record` and the trailing assistant turn is lifted out as the
    `response` side — so the bench works on whatever `pack_dataset.py` emits
    (currently the flat intermediate) without a separate test split."""
    if record.get("format") == "eliza_native_v1":
        return record
    formatted = format_record(record)
    if not formatted or not isinstance(formatted.get("messages"), list):
        return None
    messages = list(formatted["messages"])
    if not messages or messages[-1].get("role") != "assistant":
        return None
    assistant = messages[-1]
    req_messages = messages[:-1]
    system = ""
    if req_messages and req_messages[0].get("role") == "system":
        system = req_messages[0].get("content") or ""
        req_messages = req_messages[1:]
    tool_calls: list[dict[str, Any]] = []
    for tc in assistant.get("tool_calls") or []:
        fn = tc.get("function") if isinstance(tc.get("function"), dict) else {}
        args = fn.get("arguments")
        if isinstance(args, str):
            try:
                args = json.loads(args)
            except json.JSONDecodeError:
                args = {}
        tool_calls.append({"toolName": fn.get("name") or tc.get("name"),
                           "args": args if isinstance(args, dict) else {}})
    content = assistant.get("content") or ""
    if not tool_calls and isinstance(content, str):
        tool_calls = extract_tool_calls_from_text(content)
    return {
        "format": "eliza_native_v1",
        "request": {"system": system, "messages": req_messages,
                    "tools": formatted.get("tools")},
        "response": {"text": content if isinstance(content, str) else "",
                     "toolCalls": tool_calls},
        "metadata": record.get("metadata") if isinstance(record.get("metadata"), dict)
                    else {"task_type": record.get("purpose") or "response"},
    }


def load_records(path: Path, max_per_bucket: int) -> dict[str, list[dict[str, Any]]]:
    buckets: dict[str, list[dict[str, Any]]] = defaultdict(list)
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            try:
                record = json.loads(line)
            except json.JSONDecodeError:
                continue
            record = _to_native(record)
            if record is None:
                continue
            bucket = classify(record)
            if len(buckets[bucket]) < max_per_bucket:
                buckets[bucket].append(record)
    return buckets


def _coerce_args_to_dict(args: Any) -> dict[str, Any]:
    if isinstance(args, str):
        try:
            parsed = json.loads(args)
        except json.JSONDecodeError:
            return {}
        return parsed if isinstance(parsed, dict) else {}
    return args if isinstance(args, dict) else {}


def _to_openai_tool_call(raw: dict[str, Any], index: int) -> dict[str, Any] | None:
    """Normalize any tool-call shape into the OpenAI-nested shape the Gemma-4
    chat template consumes: `{"id", "type": "function", "function": {"name",
    "arguments"}}` with `arguments` as a dict. Accepts OpenAI-nested calls,
    flat `{name, arguments}` / `{toolName, args}`, and Vercel AI SDK
    `{type: "tool-call", toolName, input, toolCallId}` parts."""
    function = raw.get("function") if isinstance(raw.get("function"), dict) else {}
    name = raw.get("toolName") or raw.get("name") or function.get("name")
    if not isinstance(name, str) or not name.strip():
        return None
    args: Any = None
    for key in ("input", "args", "arguments"):
        if key in raw:
            args = raw[key]
            break
    if args is None:
        args = function.get("arguments")
    call_id = raw.get("toolCallId") or raw.get("id") or function.get("id") or f"call_{index}"
    return {
        "id": str(call_id),
        "type": "function",
        "function": {"name": name, "arguments": _coerce_args_to_dict(args)},
    }


def _tool_result_text(parts: list[Any]) -> str:
    """Flatten a Vercel `tool-result` content-parts array to a text string the
    chat template renders as the tool response body."""
    bits: list[str] = []
    for part in parts:
        if not isinstance(part, dict):
            continue
        output = part.get("output")
        if isinstance(output, dict):
            value = output.get("value")
            bits.append(value if isinstance(value, str)
                       else json.dumps(output, ensure_ascii=False))
        elif isinstance(output, str):
            bits.append(output)
        elif part.get("type") == "text" and isinstance(part.get("text"), str):
            bits.append(part["text"])
        elif "result" in part:
            result = part["result"]
            bits.append(result if isinstance(result, str)
                        else json.dumps(result, ensure_ascii=False))
    return "\n".join(b for b in bits if b)


def _normalize_message_for_template(msg: Any) -> dict[str, Any] | None:
    """Coerce one request message into the OpenAI/ChatML shape the Gemma-4
    chat template expects. Converts Vercel AI SDK content-parts (assistant
    `tool-call` parts → `tool_calls`, tool `tool-result` parts → string
    content + `tool_call_id`) and normalizes any pre-existing tool_calls."""
    if not isinstance(msg, dict):
        return None
    role = msg.get("role")
    if role == "model":
        role = "assistant"
    if role not in ("system", "user", "assistant", "tool"):
        return None
    content = msg.get("content")

    raw_tool_calls = msg.get("tool_calls")
    if raw_tool_calls is None:
        raw_tool_calls = msg.get("toolCalls")
    tool_calls: list[dict[str, Any]] = []
    if isinstance(raw_tool_calls, list):
        tool_calls = [
            call for i, raw in enumerate(raw_tool_calls)
            if (call := _to_openai_tool_call(raw, i)) is not None
        ]

    if role == "assistant" and isinstance(content, list):
        text_bits: list[str] = []
        for part in content:
            if not isinstance(part, dict):
                continue
            ptype = part.get("type")
            if ptype in ("tool-call", "tool_call"):
                call = _to_openai_tool_call(part, len(tool_calls))
                if call is not None:
                    tool_calls.append(call)
            elif ptype == "text" and isinstance(part.get("text"), str):
                text_bits.append(part["text"])
        out: dict[str, Any] = {"role": "assistant", "content": "".join(text_bits)}
        if tool_calls:
            out["tool_calls"] = tool_calls
        return out

    if role == "tool" and isinstance(content, list):
        out = {"role": "tool", "content": _tool_result_text(content)}
        call_id: Any = None
        for part in content:
            if isinstance(part, dict) and part.get("toolCallId"):
                call_id = part["toolCallId"]
                break
        if call_id is None:
            call_id = msg.get("tool_call_id")
        if call_id is not None:
            out["tool_call_id"] = str(call_id)
        if isinstance(msg.get("name"), str):
            out["name"] = msg["name"]
        return out

    out = {"role": role}
    if isinstance(content, str):
        out["content"] = content
    elif content is None:
        out["content"] = ""
    else:
        out["content"] = json.dumps(content, ensure_ascii=False)
    if tool_calls:
        out["tool_calls"] = tool_calls
    if isinstance(msg.get("tool_call_id"), str):
        out["tool_call_id"] = msg["tool_call_id"]
    if isinstance(msg.get("name"), str):
        out["name"] = msg["name"]
    return out


def _normalize_tools_for_template(tools: Any) -> list[dict[str, Any]] | None:
    """Wrap flat native tool specs (`{name, description, parameters}`) into the
    OpenAI-nested `{type: "function", function: {...}}` shape the Gemma-4
    template's `format_function_declaration` macro requires."""
    if not isinstance(tools, list):
        return None
    out: list[dict[str, Any]] = []
    for tool in tools:
        if not isinstance(tool, dict):
            continue
        if isinstance(tool.get("function"), dict):
            out.append(tool)
            continue
        name = tool.get("name")
        if not isinstance(name, str) or not name.strip():
            continue
        out.append({
            "type": "function",
            "function": {
                "name": name,
                "description": tool.get("description")
                if isinstance(tool.get("description"), str) else "",
                "parameters": tool.get("parameters")
                if isinstance(tool.get("parameters"), dict)
                else {"type": "object", "properties": {}},
            },
        })
    return out or None


def render_prompt(record: dict[str, Any], tokenizer: Any) -> tuple[str, dict[str, Any]]:
    """Render the request side of a native record for generation.

    Builds the conversation directly from `record["request"]` (every record in
    a bucket is already `eliza_native_v1` after `_to_native`, so this handles
    base-ChatML-derived and true-native rows uniformly and avoids the
    double-`format_record` path that dropped the `boundary`-less base rows).
    The supervised response lives in `record["response"]`, never in
    `request.messages`, so nothing is trimmed here."""
    request = record.get("request") if isinstance(record.get("request"), dict) else {}
    src: list[Any] = []
    system = request.get("system")
    if isinstance(system, str) and system.strip():
        src.append({"role": "system", "content": system})
    raw_messages = request.get("messages")
    if isinstance(raw_messages, list):
        src.extend(raw_messages)

    messages = [
        norm for msg in src
        if (norm := _normalize_message_for_template(msg)) is not None
    ]
    if not messages:
        return "", {}

    tools = _normalize_tools_for_template(request.get("tools"))
    kwargs: dict[str, Any] = {
        "conversation": messages,
        "tokenize": False,
        "add_generation_prompt": True,
    }
    if tools is not None:
        kwargs["tools"] = tools
    try:
        prompt = tokenizer.apply_chat_template(**kwargs)
    except TypeError:
        kwargs.pop("tools", None)
        prompt = tokenizer.apply_chat_template(**kwargs)
    return prompt, {"messages": messages, "tools": tools}


def score_tool_calls(
    predicted_text: str,
    expected_calls: list[dict[str, Any]],
) -> tuple[bool, bool, bool, dict[str, bool]]:
    predicted_calls = extract_tool_calls_from_text(predicted_text)
    parse_error = not predicted_calls
    predicted_names = [call["name"] for call in predicted_calls]
    expected_names = [call["name"] for call in expected_calls]
    expected_arg_keys = [
        sorted(call.get("arguments", {}).keys()) for call in expected_calls
    ]
    predicted_arg_keys = [
        sorted(call.get("arguments", {}).keys()) for call in predicted_calls
    ]
    fields = {
        "tool_count": len(predicted_calls) == len(expected_calls),
        "tool_names": predicted_names == expected_names,
        "argument_keys": predicted_arg_keys == expected_arg_keys,
    }
    return bool(predicted_calls), all(fields.values()), parse_error, fields


def score_json(
    predicted_text: str,
    expected_text: str,
) -> tuple[bool, bool, bool, dict[str, bool]]:
    predicted = _parse_json_object(predicted_text)
    expected = _parse_json_object(expected_text)
    if expected is None:
        ok = bool(predicted_text.strip())
        return ok, ok, False, {"nonempty": ok}
    if predicted is None:
        return False, False, True, {"json_parse": False}

    pred_handler = predicted.get("messageHandler")
    exp_handler = expected.get("messageHandler")
    if isinstance(pred_handler, dict) or isinstance(exp_handler, dict):
        pred_handler_dict = pred_handler if isinstance(pred_handler, dict) else {}
        exp_handler_dict = exp_handler if isinstance(exp_handler, dict) else {}
        pred_action = str(pred_handler_dict.get("action", "")).upper()
        exp_action = str(exp_handler_dict.get("action", "")).upper()
        pred_contexts = pred_handler_dict.get("contexts")
        exp_contexts = exp_handler_dict.get("contexts")
        fields = {
            "json_parse": True,
            "action": pred_action == exp_action,
            "contexts": isinstance(pred_contexts, list)
            and isinstance(exp_contexts, list)
            and pred_contexts == exp_contexts,
        }
        return True, fields["action"] and fields["contexts"], False, fields

    expected_keys = set(expected.keys())
    predicted_keys = set(predicted.keys())
    fields = {
        "json_parse": True,
        "top_level_keys": expected_keys.issubset(predicted_keys),
    }
    return True, fields["top_level_keys"], False, fields


def generate(
    model: Any, tokenizer: Any, prompt: str, *, max_new_tokens: int,
) -> tuple[str, int, int, float]:
    """Returns (decoded_text, n_prompt_tokens, n_gen_tokens, generate_seconds)."""
    import torch

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    prompt_len = int(inputs["input_ids"].shape[-1])
    with torch.no_grad():
        t0 = time.perf_counter()
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
        dt = time.perf_counter() - t0
    generated = output[0][prompt_len:]
    return (tokenizer.decode(generated, skip_special_tokens=False),
            prompt_len, int(generated.shape[0]), dt)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--test-file", required=True)
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--max-per-bucket", type=int, default=200)
    parser.add_argument("--max-new-tokens", type=int, default=512)
    parser.add_argument(
        "--device", default="auto", choices=("auto", "cuda", "mps", "cpu"),
        help="Inference device. 'auto' prefers cuda, then Apple-silicon mps, "
             "then cpu. mps runs the bf16 checkpoint on the M-series GPU.",
    )
    parser.add_argument(
        "--dtype", default="auto", choices=("auto", "bfloat16", "float16", "float32"),
        help="Model weight dtype. 'auto' = bfloat16 on cuda/mps (Gemma's native "
             "dtype), float32 on cpu.",
    )
    args = parser.parse_args()

    buckets = load_records(Path(args.test_file), args.max_per_bucket)
    if not buckets:
        raise SystemExit(f"no usable records found in {args.test_file}")

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    if args.device != "auto":
        device = args.device
    elif torch.cuda.is_available():
        device = "cuda"
    elif getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"
    if args.dtype != "auto":
        dtype = getattr(torch, args.dtype)
    else:
        dtype = torch.float32 if device == "cpu" else torch.bfloat16
    log.info("device=%s dtype=%s", device, dtype)
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    # Base Gemma-4 tokenizers ship no chat_template; this bench renders prompts
    # via apply_chat_template, so borrow the template from the -it instruct
    # variant when the base has none (mirrors train_local.py).
    if not getattr(tokenizer, "chat_template", None):
        try:
            _src = AutoTokenizer.from_pretrained(
                f"{args.model}-it", trust_remote_code=True
            )
            if getattr(_src, "chat_template", None):
                tokenizer.chat_template = _src.chat_template
        except Exception:  # noqa: BLE001
            pass
    model_kwargs: dict[str, Any] = {
        "torch_dtype": dtype,
        "trust_remote_code": True,
        "low_cpu_mem_usage": True,
        "attn_implementation": select_attn_impl(device),
    }
    if device == "cuda":
        model_kwargs["device_map"] = "auto"
    model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)
    if device != "cuda":
        model.to(device)
    model.eval()

    results: dict[str, BucketResult] = {
        bucket: BucketResult(bucket) for bucket in buckets
    }
    for bucket, records in buckets.items():
        log.info("bucket=%s n=%d", bucket, len(records))
        for record in records:
            prompt, _formatted = render_prompt(record, tokenizer)
            if not prompt:
                continue
            predicted, n_prompt_tok, n_gen_tok, gen_dt = generate(
                model,
                tokenizer,
                prompt,
                max_new_tokens=args.max_new_tokens,
            )
            expected_calls = expected_tool_calls(record)
            if expected_calls:
                ok_structure, ok_content, parse_error, fields = score_tool_calls(
                    predicted,
                    expected_calls,
                )
            else:
                ok_structure, ok_content, parse_error, fields = score_json(
                    predicted,
                    response_text(record),
                )
            results[bucket].record(
                ok_structure=ok_structure,
                ok_content=ok_content,
                parse_error=parse_error,
                fields=fields,
                gen_dt=gen_dt,
                n_prompt_tokens=n_prompt_tok,
                n_gen_tokens=n_gen_tok,
                failed=None
                if ok_content
                else {
                    "trajectoryId": record.get("trajectoryId"),
                    "callId": record.get("callId"),
                    "expected": record.get("response"),
                    "predicted": predicted[:2000],
                },
            )

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    total_gen_seconds = sum(r.gen_seconds for r in results.values())
    total_prompt_tokens = sum(r.prompt_tokens for r in results.values())
    total_gen_tokens = sum(r.gen_tokens for r in results.values())
    summary = {
        "model": args.model,
        "test_file": args.test_file,
        # prompt_tps / gen_tps are over the generate()-only wallclock
        # (prefill+decode, not separated by model.generate).
        "prompt_tps": round(total_prompt_tokens / max(total_gen_seconds, 1e-6), 2),
        "gen_tps": round(total_gen_tokens / max(total_gen_seconds, 1e-6), 2),
        "gen_seconds": round(total_gen_seconds, 2),
        "buckets": {key: result.to_dict() for key, result in results.items()},
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())