Release extreme 48-block-pruned experimental edition with benchmark and quality warnings
7ae0b4d verified | """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() | |