import argparse import math import time import jax import jax.numpy as jnp import numpy as np import optax from .data import get_tokenizer, load_tool_calls, load_prepared_data, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_DEC_LEN, DEFAULT_MAX_GEN_LEN from .model import ( EncoderDecoderTransformer, TransformerConfig, make_causal_mask, make_causal_packing_mask, make_padding_mask, make_packing_mask, make_cross_packing_mask, ) from .run import load_checkpoint, _get_decode_fn def compute_perplexity_packed(model, params, packed_data, batch_size): """Compute perplexity from pre-packed data with segment IDs.""" enc = packed_data["packed_enc"] dec_in = packed_data["packed_dec_in"] dec_tgt = packed_data["packed_dec_tgt"] enc_seg = packed_data["packed_enc_seg"] dec_seg = packed_data["packed_dec_seg"] total_loss = 0.0 total_tokens = 0 n = len(enc) for i in range(0, n - batch_size + 1, batch_size): src = jnp.asarray(enc[i:i+batch_size], dtype=jnp.int32) tgt_in = jnp.asarray(dec_in[i:i+batch_size], dtype=jnp.int32) tgt_out = jnp.asarray(dec_tgt[i:i+batch_size], dtype=jnp.int32) es = jnp.asarray(enc_seg[i:i+batch_size], dtype=jnp.int32) ds = jnp.asarray(dec_seg[i:i+batch_size], dtype=jnp.int32) src_mask = make_packing_mask(es) tgt_mask = make_causal_packing_mask(ds) cross_mask = make_cross_packing_mask(es, ds) logits = model.apply( {"params": params}, src, tgt_in, src_mask=src_mask, tgt_mask=tgt_mask, cross_mask=cross_mask, ) loss = optax.softmax_cross_entropy_with_integer_labels(logits.astype(jnp.float32), tgt_out) padding_mask = (ds > 0).astype(jnp.float32) total_loss += float(jnp.sum(loss * padding_mask)) total_tokens += int(jnp.sum(padding_mask)) avg_nll = total_loss / max(total_tokens, 1) return math.exp(min(avg_nll, 20)) def measure_throughput(model, params, tokenizer, num_runs=10, prompt='What is the weather?', max_gen_len=64): enc_tokens = tokenizer.encode(prompt) enc_input = jnp.array([enc_tokens]) pad_id = tokenizer.pad_token_id eos_id = tokenizer.eos_token_id src_mask = make_padding_mask(enc_input, pad_id) decode_fn = _get_decode_fn(model, max_gen_len) # Warmup JIT encoder_out, enc_mask = model.apply( {"params": params}, enc_input, src_mask=src_mask, method="encode" ) dec_buffer = jnp.full((1, max_gen_len), pad_id, dtype=jnp.int32) dec_buffer = dec_buffer.at[0, 0].set(eos_id) decode_fn(params, dec_buffer, encoder_out, enc_mask) tokens_generated = [] latencies = [] for run in range(num_runs): rng = jax.random.PRNGKey(run) dec_buffer = jnp.full((1, max_gen_len), pad_id, dtype=jnp.int32) dec_buffer = dec_buffer.at[0, 0].set(eos_id) start = time.perf_counter() encoder_out, enc_mask = model.apply( {"params": params}, enc_input, src_mask=src_mask, method="encode" ) logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) num_tokens = 0 for i in range(max_gen_len - 1): rng, sample_rng = jax.random.split(rng) next_token = jax.random.categorical(sample_rng, logits[0, i]).item() if next_token == eos_id: break num_tokens += 1 dec_buffer = dec_buffer.at[0, i + 1].set(next_token) logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) elapsed = time.perf_counter() - start tokens_generated.append(num_tokens) latencies.append(elapsed) avg_tokens = np.mean(tokens_generated) avg_latency = np.mean(latencies) tokens_per_sec = sum(tokens_generated) / sum(latencies) return { "avg_tokens_generated": avg_tokens, "avg_latency_s": avg_latency, "tokens_per_second": tokens_per_sec, } def compute_repetition_rate(texts): bigram_rep_rates = [] for text in texts: words = text.lower().split() if len(words) < 2: bigram_rep_rates.append(0.0) continue bigrams = [(words[i], words[i + 1]) for i in range(len(words) - 1)] unique = len(set(bigrams)) bigram_rep_rates.append(1.0 - unique / len(bigrams)) return float(np.mean(bigram_rep_rates)) def benchmark_generation_quality(model, params, tokenizer, prompts, max_gen_len=128, temperature=0.8): from .run import generate generations = [] for i, prompt in enumerate(prompts): text = generate(model, params, tokenizer, prompt, max_gen_len, temperature, seed=i, stream=False) generations.append(text) lengths = [len(tokenizer.encode(t)) for t in generations] rep_rate = compute_repetition_rate(generations) return { "avg_generation_length": float(np.mean(lengths)), "min_generation_length": int(np.min(lengths)), "max_generation_length": int(np.max(lengths)), "bigram_repetition_rate": rep_rate, "generations": list(zip(prompts, generations)), } def compute_wer(hypotheses, references): """Compute word error rate using edit distance.""" total_edits = 0 total_ref_words = 0 for hyp, ref in zip(hypotheses, references): hyp_words = hyp.lower().split() ref_words = ref.lower().split() n = len(ref_words) m = len(hyp_words) # DP edit distance d = [[0] * (m + 1) for _ in range(n + 1)] for i in range(n + 1): d[i][0] = i for j in range(m + 1): d[0][j] = j for i in range(1, n + 1): for j in range(1, m + 1): if ref_words[i - 1] == hyp_words[j - 1]: d[i][j] = d[i - 1][j - 1] else: d[i][j] = 1 + min(d[i - 1][j], d[i][j - 1], d[i - 1][j - 1]) total_edits += d[n][m] total_ref_words += n return total_edits / max(total_ref_words, 1) def benchmark_tool_calls(model, params, tokenizer, num_samples=200, max_gen_len=DEFAULT_MAX_GEN_LEN, max_enc_len=DEFAULT_MAX_ENC_LEN, constrained=True, ds=None): """Generate tool-call predictions and compute structured metrics.""" import json from .run import generate_batch, normalize_tools, restore_tool_names from .data import load_tool_calls, to_snake_case if ds is None: ds = load_tool_calls("validation", max_samples=num_samples) queries = [ex["query"] for ex in ds] tools_list = [ex["tools"] for ex in ds] all_preds = generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=max_gen_len, max_enc_len=max_enc_len, normalize=True, constrained=constrained) total = 0 exact_match = 0 tp_names = 0 fp_names = 0 fn_names = 0 tp_calls = 0 fp_calls = 0 fn_calls = 0 json_parse_errors = 0 empty_ref = 0 empty_pred = 0 args_correct = 0 args_total = 0 halluc_params = 0 total_pred_params = 0 missing_params = 0 total_ref_params = 0 correct_values = 0 matched_params = 0 samples = [] failures = [] # (query, tools, ref, pred, reasons) def _normalize_value(v): """Normalize a value for fuzzy comparison.""" if isinstance(v, str): # Try parsing as number to handle "74.006" vs "74.0060" try: f = float(v) # Preserve sign, normalize trailing zeros v = str(f) except ValueError: pass # Normalize time-like strings: strip leading "at ", lowercase s = v.strip().lower() if s.startswith("at "): s = s[3:].strip() # "today at 21:00" → "21:00", "tonight" is harder but normalize what we can if s.startswith("today at "): s = s[len("today at "):].strip() return s if isinstance(v, float): return str(v) return v def _normalize_args(args): """Normalize all argument values for comparison.""" if not isinstance(args, dict): return args return {k: _normalize_value(v) for k, v in args.items()} def call_key(c): if not isinstance(c, dict): return None norm_args = _normalize_args(c.get("arguments", {})) return json.dumps({"name": c.get("name"), "arguments": norm_args}, sort_keys=True) for i, (ex, pred_text) in enumerate(zip(ds, all_preds)): ref_text = ex["answers"] pred_text = pred_text.strip() try: ref_calls = json.loads(ref_text) except (json.JSONDecodeError, TypeError): ref_calls = [] # Normalize ref tool names to snake_case for consistent comparison for rc in ref_calls: if isinstance(rc, dict) and "name" in rc: rc["name"] = to_snake_case(rc["name"]) # Normalize pred tool names to snake_case too (restore_tool_names in # generate_batch maps back to originals, but we want snake_case comparison) try: pred_calls_raw = json.loads(pred_text) if isinstance(pred_calls_raw, list): for pc in pred_calls_raw: if isinstance(pc, dict) and "name" in pc: pc["name"] = to_snake_case(pc["name"]) elif isinstance(pred_calls_raw, dict) and "name" in pred_calls_raw: pred_calls_raw["name"] = to_snake_case(pred_calls_raw["name"]) pred_text = json.dumps(pred_calls_raw) except (json.JSONDecodeError, TypeError): pass # Also normalize tool definitions for param validation try: tool_defs_raw = json.loads(ex["tools"]) for td in tool_defs_raw: if isinstance(td, dict) and "name" in td: td["name"] = to_snake_case(td["name"]) tools_normalized = json.dumps(tool_defs_raw) except (json.JSONDecodeError, TypeError): tools_normalized = ex["tools"] try: pred_calls = json.loads(pred_text) if not isinstance(pred_calls, list): pred_calls = [pred_calls] if isinstance(pred_calls, dict) else [] except (json.JSONDecodeError, TypeError): json_parse_errors += 1 pred_calls = [] total += 1 if i < 10: samples.append((ex["query"][:80], ref_text[:120], pred_text[:120])) ref_is_empty = ref_text.strip() in ("", "[]") pred_is_empty = pred_text.strip() in ("", "[]") if ref_is_empty: empty_ref += 1 if pred_is_empty: empty_pred += 1 if ref_is_empty and pred_is_empty: exact_match += 1 elif not ref_is_empty and not pred_is_empty: # Order-independent comparison: sort calls by their canonical key ref_sorted = sorted([call_key(c) for c in ref_calls if call_key(c)]) pred_sorted = sorted([call_key(c) for c in pred_calls if call_key(c)]) if ref_sorted == pred_sorted and len(ref_sorted) == len(ref_calls) and len(pred_sorted) == len(pred_calls): exact_match += 1 ref_name_set = {c["name"] for c in ref_calls if isinstance(c, dict) and "name" in c} pred_name_set = {c["name"] for c in pred_calls if isinstance(c, dict) and "name" in c} tp_names += len(pred_name_set & ref_name_set) fp_names += len(pred_name_set - ref_name_set) fn_names += len(ref_name_set - pred_name_set) ref_keys = {call_key(c) for c in ref_calls} - {None} pred_keys = {call_key(c) for c in pred_calls} - {None} tp_calls += len(pred_keys & ref_keys) fp_calls += len(pred_keys - ref_keys) fn_calls += len(ref_keys - pred_keys) ref_by_name = {} for c in ref_calls: if isinstance(c, dict) and "name" in c: ref_by_name.setdefault(c["name"], []).append(c.get("arguments", {})) for c in pred_calls: if isinstance(c, dict) and "name" in c and c["name"] in ref_by_name: args_total += 1 pa = json.dumps(_normalize_args(c.get("arguments", {})), sort_keys=True) if any(pa == json.dumps(_normalize_args(ra), sort_keys=True) for ra in ref_by_name[c["name"]]): args_correct += 1 try: tool_defs = json.loads(tools_normalized) tool_param_map = {t["name"]: set((t.get("parameters") or {}).keys()) for t in tool_defs if isinstance(t, dict) and "name" in t} except (json.JSONDecodeError, TypeError): tool_param_map = {} for c in pred_calls: if not isinstance(c, dict) or "name" not in c: continue cname = c["name"] if cname not in tool_param_map: continue schema_keys = tool_param_map[cname] p_keys = set((c.get("arguments") or {}).keys()) total_pred_params += len(p_keys) halluc_params += len(p_keys - schema_keys) if cname in ref_by_name: ref_args = ref_by_name[cname][0] r_keys = set((ref_args if isinstance(ref_args, dict) else {}).keys()) total_ref_params += len(r_keys) missing_params += len(r_keys - p_keys) m_keys = p_keys & r_keys matched_params += len(m_keys) for k in m_keys: pv_norm = _normalize_value(c.get("arguments", {})[k]) rv_norm = _normalize_value(ref_args[k]) if json.dumps(pv_norm, sort_keys=True) == json.dumps(rv_norm, sort_keys=True): correct_values += 1 # Diagnose failures for this example is_exact = (ref_is_empty and pred_is_empty) or ( not ref_is_empty and not pred_is_empty and sorted([call_key(c) for c in ref_calls if call_key(c)]) == sorted([call_key(c) for c in pred_calls if call_key(c)]) and len(ref_calls) == len(pred_calls) ) if not is_exact and len(failures) < 50: reasons = [] if json_parse_errors > 0 and i == total - 1: # Check if this specific example had a parse error try: json.loads(pred_text) except (json.JSONDecodeError, TypeError): reasons.append("json_parse_error") ex_fp_names = pred_name_set - ref_name_set ex_fn_names = ref_name_set - pred_name_set if ex_fp_names: reasons.append(f"wrong_tools:{','.join(sorted(ex_fp_names))}") if ex_fn_names: reasons.append(f"missing_tools:{','.join(sorted(ex_fn_names))}") if not ex_fp_names and not ex_fn_names and pred_name_set: # Right tools but wrong args — find which values differ for c in pred_calls: if not isinstance(c, dict) or "name" not in c: continue cname = c["name"] if cname not in ref_by_name: continue pred_args = c.get("arguments", {}) ref_args = ref_by_name[cname][0] for k in set(pred_args.keys()) | set(ref_args.keys()): pv = pred_args.get(k, "") rv = ref_args.get(k, "") pv_n = _normalize_value(pv) rv_n = _normalize_value(rv) if json.dumps(pv_n, sort_keys=True) != json.dumps(rv_n, sort_keys=True): reasons.append(f"value_mismatch:{cname}.{k}={json.dumps(pv)[:60]}!={json.dumps(rv)[:60]}") if ref_is_empty and not pred_is_empty: reasons.append("false_positive:should_be_empty") if not ref_is_empty and pred_is_empty: reasons.append("false_negative:predicted_empty") if not reasons: reasons.append("unknown") failures.append({ "query": ex["query"][:200], "ref": ref_text[:300], "pred": pred_text[:300], "reasons": reasons, }) name_precision = tp_names / max(tp_names + fp_names, 1) name_recall = tp_names / max(tp_names + fn_names, 1) name_f1 = 2 * name_precision * name_recall / max(name_precision + name_recall, 1e-9) call_precision = tp_calls / max(tp_calls + fp_calls, 1) call_recall = tp_calls / max(tp_calls + fn_calls, 1) call_f1 = 2 * call_precision * call_recall / max(call_precision + call_recall, 1e-9) return { "num_samples": total, "exact_match": exact_match / max(total, 1), "json_parse_rate": 1.0 - json_parse_errors / max(total, 1), "name_precision": name_precision, "name_recall": name_recall, "name_f1": name_f1, "args_acc": args_correct / max(args_total, 1), "param_haluc": halluc_params / max(total_pred_params, 1), "param_miss": missing_params / max(total_ref_params, 1), "value_acc": correct_values / max(matched_params, 1), "call_precision": call_precision, "call_recall": call_recall, "call_f1": call_f1, "empty_ref_pct": empty_ref / max(total, 1), "empty_pred_pct": empty_pred / max(total, 1), "samples": samples, "failures": failures, } def benchmark_retrieval(model, params, tokenizer, num_samples=500, max_len=256, ks=(1, 2, 3, 4, 5), ds=None): """Benchmark contrastive retrieval: Recall@k and MRR over validation set. For each query, ranks all tools from that example by cosine similarity and checks if the positive (called) tools appear in top-k. """ import json from .run import encode_for_retrieval if ds is None: ds = load_tool_calls("validation", max_samples=num_samples) elif num_samples and len(ds) > num_samples: ds = ds.select(range(num_samples)) queries = [] all_tool_strs = [] tool_groups = [] # (start_idx, count, set_of_positive_indices_within_group) for ex in ds: query = ex["query"] try: tools = json.loads(ex["tools"]) except (ValueError, TypeError): continue if not isinstance(tools, list) or len(tools) == 0: continue try: calls = json.loads(ex["answers"]) except (ValueError, TypeError): calls = [] pos_names = set() if isinstance(calls, list): for c in calls: if isinstance(c, dict) and "name" in c: pos_names.add(c["name"]) if not pos_names: continue start = len(all_tool_strs) pos_indices = set() for j, tool in enumerate(tools): if not isinstance(tool, dict): continue all_tool_strs.append(json.dumps(tool, separators=(",", ":"))) if tool.get("name", "") in pos_names: pos_indices.add(len(all_tool_strs) - 1 - start) if pos_indices: queries.append(query) tool_groups.append((start, len(all_tool_strs) - start, pos_indices)) if not queries: return {"recall@k": {}, "mrr": 0.0, "num_queries": 0} # Encode all queries and tools q_embs = encode_for_retrieval(model, params, tokenizer, queries, max_len=max_len) t_embs = encode_for_retrieval(model, params, tokenizer, all_tool_strs, max_len=max_len) max_k = max(ks) recall = {k: 0 for k in ks} mrr_sum = 0.0 for i, (start, count, pos_set) in enumerate(tool_groups): group_embs = t_embs[start:start + count] # (count, dim) scores = q_embs[i:i+1] @ group_embs.T # (1, count) ranked = np.argsort(-scores[0]) # MRR: reciprocal rank of first positive for rank, idx in enumerate(ranked): if idx in pos_set: mrr_sum += 1.0 / (rank + 1) break # Recall@k for k in ks: top_k_set = set(ranked[:k].tolist()) if top_k_set & pos_set: recall[k] += 1 n = len(queries) return { "recall@k": {k: recall[k] / n for k in ks}, "mrr": mrr_sum / n, "num_queries": n, } def main(args): import json as _json_mod params, config = load_checkpoint(args.checkpoint) model = EncoderDecoderTransformer(config) tokenizer = get_tokenizer() param_count = sum(x.size for x in jax.tree.leaves(params)) print(f"\ncheckpoint: {args.checkpoint}") print(f"parameters: {param_count:,}") print(f"config: d={config.d_model}, heads={config.num_heads}, layers={config.num_encoder_layers}/{config.num_decoder_layers}") # --- Perplexity from pre-packed data (no re-tokenization) --- print(f"\nevaluating perplexity (pre-packed val data)...") try: val_data = load_prepared_data("val", mmap=True) n_bins = len(val_data["packed_enc"]) ppl = compute_perplexity_packed(model, params, val_data, args.batch_size) print(f" perplexity: {ppl:.2f} ({n_bins} packed bins)") except FileNotFoundError: print(" Skipped — no pre-packed val data. Run 'needle tokenize' first.") ppl = None # --- Tool-call accuracy with stratified sampling (matches training eval) --- tc_samples = getattr(args, "tool_call_samples", 200) use_constrained = not getattr(args, "no_constrained", False) if tc_samples > 0: val_ds = load_tool_calls("validation") rng = np.random.RandomState(42) def _classify(ex): try: answers = _json_mod.loads(ex["answers"]) except (ValueError, TypeError): return "empty" if not answers or not isinstance(answers, list): return "empty" return "single" if len(answers) == 1 else "multi" pool_names = ["single", "multi"] pool_buckets = {name: {t: [] for t in range(11)} for name in pool_names} for k in range(len(val_ds)): ex = val_ds[k] try: nc = min(len(_json_mod.loads(ex["tools"])), 10) except (ValueError, TypeError): nc = 0 ct = _classify(ex) if ct != "empty": pool_buckets[ct][nc].append(k) per_bucket = max(1, tc_samples // 14) # ~7 tool-count buckets * 2 pools for name in pool_names: label = "Single-Call" if name == "single" else "Multi-Call" pool = [] for t in range(11): b = np.array(pool_buckets[name][t]) if len(b) > 0: rng.shuffle(b) pool.extend(b[:per_bucket].tolist()) rng.shuffle(pool) examples = [val_ds[k] for k in pool] if not examples: continue print(f"\n evaluating {label} ({len(examples)} samples, constrained={use_constrained})...") tc = benchmark_tool_calls( model, params, tokenizer, num_samples=len(examples), max_gen_len=args.max_gen_len, max_enc_len=args.max_enc_len, constrained=use_constrained, ds=examples, ) print(f" ─── {label} ({len(examples)} samples) ───") print(f" JSON parse {tc['json_parse_rate']:>10.1%}") print(f" Name F1 {tc['name_f1']:>10.1%}") print(f" Param haluc {tc['param_haluc']:>10.1%}") print(f" Param miss {tc['param_miss']:>10.1%}") print(f" Value acc {tc['value_acc']:>10.1%}") print(f" Args acc {tc['args_acc']:>10.1%}") print(f" Call F1 {tc['call_f1']:>10.1%}") print(f" Exact match {tc['exact_match']:>10.1%}") if tc["samples"]: print(f" samples:") for query, ref, pred in tc["samples"][:3]: print(f" Q: {query}") print(f" R: {ref}") print(f" P: {pred}") print() if tc.get("failures"): print(f" ─── Failures ({len(tc['failures'])} captured) ───") for j, fail in enumerate(tc["failures"][:15]): print(f" [{j+1}] Q: {fail['query'][:120]}") print(f" Ref: {fail['ref'][:200]}") print(f" Pred: {fail['pred'][:200]}") print(f" Why: {', '.join(fail['reasons'])}") print() print(f"\nmeasuring throughput ({args.throughput_runs} runs)...") throughput = measure_throughput(model, params, tokenizer, num_runs=args.throughput_runs) print(f"avg tokens: {throughput['avg_tokens_generated']:.1f}") print(f"avg latency: {throughput['avg_latency_s']:.3f}s") print(f"tokens/sec: {throughput['tokens_per_second']:.1f}") def parse_args(): parser = argparse.ArgumentParser(description="Benchmark transformer on generation tasks") parser.add_argument("--checkpoint", type=str, default="checkpoints/checkpoint_epoch3.pkl") parser.add_argument("--batch-size", type=int, default=32) parser.add_argument("--max-eval-samples", type=int, default=1000) parser.add_argument("--max-enc-len", type=int, default=DEFAULT_MAX_ENC_LEN) parser.add_argument("--max-dec-len", type=int, default=DEFAULT_MAX_DEC_LEN) parser.add_argument("--max-gen-len", type=int, default=DEFAULT_MAX_GEN_LEN) parser.add_argument("--throughput-runs", type=int, default=10) parser.add_argument("--tool-call-samples", type=int, default=200, help="Samples for tool-call accuracy eval (default: 200)") return parser.parse_args() if __name__ == "__main__": main(parse_args())