Evaluate a model

import json
import random

SEED = 20261101
HOLDOUT_FRACTION = 0.1

with open("data.jsonl", encoding="utf-8") as source:
    rows = [json.loads(line) for line in source]
groups = sorted({row["group_key"] for row in rows})
if len(groups) < 2:
    raise ValueError("Use at least two groups so both splits contain records")
random.Random(SEED).shuffle(groups)
held_out = set(groups[: max(1, int(len(groups) * HOLDOUT_FRACTION))])

with (
    open("train.jsonl", "w", encoding="utf-8") as train,
    open("holdout.jsonl", "w", encoding="utf-8") as holdout,
):
    for row in rows:
        (holdout if row["group_key"] in held_out else train).write(
            json.dumps(row) + "\n"
        )

held_out_rows = sum(row["group_key"] in held_out for row in rows)
print(f"train.jsonl: {len(rows) - held_out_rows} rows")
print(f"holdout.jsonl: {held_out_rows} rows")

train.jsonl: 3 rows
holdout.jsonl: 1 rows

import json


def tool_calls(message):
    """Return sorted (name, arguments) pairs, or None if any call's arguments don't parse."""
    calls = []
    for call in message.get("tool_calls") or []:
        arguments = call["function"]["arguments"]
        # API replies send JSON text; datasets often store objects.
        if isinstance(arguments, str):
            try:
                arguments = json.loads(arguments)
            except json.JSONDecodeError:
                return None
        calls.append((call["function"]["name"], json.dumps(arguments, sort_keys=True)))
    return sorted(calls)


def score(reply, finish_reason, reference):
    got, want = tool_calls(reply), tool_calls(reference)
    if want is None:
        raise ValueError(
            "the reference's tool call arguments don't parse; fix the held-out row"
        )
    return {
        "decision": got is not None and bool(got) == bool(want),
        "tools": got is not None and [n for n, _ in got] == [n for n, _ in want],
        "arguments": got == want,
        "well_formed": got is not None and finish_reason in ("stop", "tool_calls"),
    }

import os
from typing import Literal

from openai import OpenAI
from pydantic import BaseModel

client = OpenAI(
    base_url="https://inference.baseten.co/v1", api_key=os.environ["BASETEN_API_KEY"]
)

RULE = (
    "The reply answers the user's request directly and completely, "
    "without calling or describing a tool."
)


class Verdict(BaseModel):
    reasoning: str
    evidence: str
    verdict: Literal["pass", "fail"]


def judge(request: str, reply: str) -> bool:
    response = client.beta.chat.completions.parse(
        model="moonshotai/Kimi-K3",
        messages=[
            {
                "role": "system",
                "content": f"Grade one rule. Rule: {RULE}\n"
                "Explain your reasoning, quote the part of the reply that decides it, "
                "then give the verdict.",
            },
            {"role": "user", "content": f"Request:\n{request}\n\nReply:\n{reply}"},
        ],
        response_format=Verdict,
        temperature=0,
    )
    verdict = response.choices[0].message.parsed
    if verdict is None:
        raise RuntimeError("The judge returned no verdict.")
    return verdict.verdict == "pass"

import json
import os

from openai import OpenAI

from checks import score

ENDPOINTS = {
    "baseline": (
        "https://model-<baseline_model_id>.api.baseten.co/deployment/<baseline_deployment_id>/sync/v1",
        "<served_model_name>",
    ),
    "candidate": (
        "https://model-<candidate_model_id>.api.baseten.co/deployment/<candidate_deployment_id>/sync/v1",
        "step-124",
    ),
}

rows = [json.loads(line) for line in open("holdout.jsonl")]
totals = {}
for label, (base_url, model) in ENDPOINTS.items():
    client = OpenAI(base_url=base_url, api_key=os.environ["BASETEN_API_KEY"])
    scores = []
    for row in rows:
        reference = next(m for m in row["messages"] if m["role"] == "assistant")
        prompt = row["messages"][: row["messages"].index(reference)]
        tools = {"tools": row["tools"]} if row.get("tools") else {}
        response = client.chat.completions.create(
            model=model, messages=prompt, temperature=0, max_tokens=1024, **tools
        )
        choice = response.choices[0]
        scores.append(
            score(choice.message.model_dump(), choice.finish_reason, reference)
        )
    totals[label] = {
        check: sum(s[check] for s in scores) / len(scores) for check in scores[0]
    }

for check in totals["baseline"]:
    print(
        f"{check:12} {totals['baseline'][check]:7.1%} {totals['candidate'][check]:7.1%}"
    )