"""Capability-focused quality gate for Bonsai's intended fast-agent workload.""" from __future__ import annotations import argparse import gc import json import math from pathlib import Path import mlx.core as mx from mlx_lm.generate import generate_step, make_sampler from transformers import AutoTokenizer from optimized_runtime import load_text_model WEB_TOOL = { "type": "function", "function": { "name": "web_search", "description": "Search the web for current information.", "parameters": { "type": "object", "properties": {"query": {"type": "string"}}, "required": ["query"], }, }, } IMAGE_TOOL = { "type": "function", "function": { "name": "generate_image", "description": "Generate an image from a detailed visual prompt.", "parameters": { "type": "object", "properties": { "prompt": {"type": "string"}, "aspect_ratio": {"type": "string"}, }, "required": ["prompt"], }, }, } TASKS = [ { "name": "web_search_tool_call", "messages": [{"role": "user", "content": "Find today's weather forecast for Detroit, Michigan."}], "tools": [WEB_TOOL], }, { "name": "current_fact_tool_call", "messages": [{"role": "user", "content": "What is the newest stable MLX release? Verify it on the web."}], "tools": [WEB_TOOL], }, { "name": "image_generator_prompt", "messages": [{"role": "user", "content": "Create a cinematic image of a tiny red spacecraft flying through a gigantic blue ice cavern, widescreen."}], "tools": [IMAGE_TOOL], }, { "name": "web_result_summary", "messages": [ { "role": "system", "content": "Summarize supplied web results accurately. Do not invent facts.", }, { "role": "user", "content": ( "Web results: [1] City transit says Route 8 closes at 9 PM Friday for bridge work. " "[2] Weekend service resumes Saturday at 6 AM with buses every 20 minutes. " "[3] Riders should use Route 12 between Central and River. Summarize in two bullets." ), }, ], }, { "name": "general_question", "messages": [{"role": "user", "content": "Explain recursion to a curious twelve-year-old in three short sentences."}], }, { "name": "clarifying_question", "messages": [ { "role": "system", "content": "If a request lacks information needed to do it well, ask one concise clarifying question.", }, {"role": "user", "content": "Make me an image for my project."}, ], "tools": [IMAGE_TOOL], }, { "name": "helpful_chat", "messages": [{"role": "user", "content": "My Python script says KeyError: 'name'. What does that mean and what should I check?"}], }, ] def encode_task(tokenizer, task): encoded = tokenizer.apply_chat_template( task["messages"], tools=task.get("tools"), tokenize=True, add_generation_prompt=True, enable_thinking=False, ) if hasattr(encoded, "input_ids"): encoded = encoded.input_ids return mx.array(encoded, dtype=mx.uint32) def make_baseline(model, tokenizer, generation_tokens, tasks): sampler = make_sampler(temp=0.0) cases = [] for task in tasks: prompt = encode_task(tokenizer, task) output, chosen_logprobs = [], [] for token, logprobs in generate_step( prompt, model, max_tokens=generation_tokens, sampler=sampler ): token = int(token) mx.eval(logprobs) output.append(token) chosen_logprobs.append(float(logprobs[token].item())) decoded = tokenizer.decode(output) cases.append( { "name": task["name"], "prompt_tokens": prompt.tolist(), "output_tokens": output, "chosen_logprobs": chosen_logprobs, "decoded": decoded, } ) print(f"baseline {task['name']}: {decoded!r}", flush=True) del prompt, output, chosen_logprobs gc.collect() mx.clear_cache() return cases def evaluate_variant(model, tokenizer, baseline): total = matches = 0 deltas = [] cases = [] for case in baseline: cache = model.make_cache() current = mx.array(case["prompt_tokens"], dtype=mx.uint32) candidate_tokens = [] case_matches = 0 case_deltas = [] for index, expected in enumerate(case["output_tokens"]): logits = model(current[None], cache=cache)[:, -1, :] logprobs = logits - mx.logsumexp(logits, keepdims=True) predicted = mx.argmax(logprobs, axis=-1) mx.eval(predicted, logprobs) predicted = int(predicted.item()) candidate_tokens.append(predicted) if predicted == expected: matches += 1 case_matches += 1 delta = float(logprobs[0, expected].item()) - case["chosen_logprobs"][index] deltas.append(delta) case_deltas.append(delta) total += 1 current = mx.array([expected], dtype=mx.uint32) cases.append( { "name": case["name"], "tokens": len(case["output_tokens"]), "top1_agreement": case_matches / len(case["output_tokens"]), "relative_perplexity": math.exp(-sum(case_deltas) / len(case_deltas)), "free_running_first_choices": tokenizer.decode(candidate_tokens), } ) print(f"evaluated {case['name']}", flush=True) mean_delta = sum(deltas) / len(deltas) return { "tokens": total, "top1_agreement": matches / total, "relative_perplexity": math.exp(-mean_delta), "cases": cases, } def main(): parser = argparse.ArgumentParser() parser.add_argument("--model", type=Path, default=Path("model")) parser.add_argument("--baseline", type=Path) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--generation-tokens", type=int, default=48) parser.add_argument( "--case", choices=[task["name"] for task in TASKS], help="run exactly one case; required for the safe local workflow", ) args = parser.parse_args() if not args.case: parser.error("--case is required; multi-case model runs are disabled for safety") mx.set_default_device(mx.gpu) mx.set_cache_limit(256_000_000) mx.set_wired_limit(mx.device_info()["max_recommended_working_set_size"]) model, config = load_text_model(args.model.resolve()) tokenizer = AutoTokenizer.from_pretrained(args.model.resolve()) if args.baseline: reference = json.loads(args.baseline.read_text()) result = { "mode": "variant", "skip_layers": config.get("active_runtime_skip_layers", []), "skip_mlp_layers": config.get("active_runtime_skip_mlp_layers", []), **evaluate_variant(model, tokenizer, reference["cases"]), } else: tasks = [task for task in TASKS if task["name"] == args.case] result = { "mode": "baseline", "generation_tokens": args.generation_tokens, "cases": make_baseline(model, tokenizer, args.generation_tokens, tasks), } args.output.write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2)) if __name__ == "__main__": main()