From 9b22a8b3449a80ab894179426267aaf14c926cf8 Mon Sep 17 00:00:00 2001 From: Joe Khawand <93840910+Joe-Khawand@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:26:34 +0200 Subject: [PATCH] Fix parallel-call eval scoring, Muon update, and assorted correctness issues Evaluation (training/eval.py) - Score calls as multisets so identical parallel calls are no longer collapsed by set-dedup (fixes call precision/recall on parallel). - Match predicted args to distinct reference instances via greedy bipartite matching instead of always ref[0], fixing value_acc and the param metrics for parallel_multiple. - Add BFCL-style category buckets (simple / multiple / parallel / parallel_multiple) plus irrelevance, relevance and false-trigger rates. - Split the scoring out into a pure score_tool_calls() so it can be exercised without a checkpoint. Training - Muon: apply Nesterov momentum to the raw gradient and orthogonalize the blended update. The previous order orthogonalized first and ran momentum on the orthogonal factors, so the applied update was a sum of orthogonal matrices (not itself orthogonal). Also add aspect-ratio update scaling. - Mask z-loss to non-pad positions in both train and pretrain; CE was already masked, so z-loss magnitude was scaling with the padding fraction. Inference (model/run.py, model/architecture.py) - decode() can project only the current position to the vocab, avoiding the full-buffer (B,T,d)@(d,V) matmul at every step. Numerically identical. - Stream the decode delta of the whole sequence rather than decoding each token in isolation, which dropped SentencePiece spaces and split multibyte byte-fallback pieces. Data (dataset/dataset.py, dataset/generate.py) - Weight all repeated tool names / argument values in the loss mask, not just the first occurrence (matters for parallel calls). - Match the boolean-polarity validator on whole words rather than substrings ("on" was matching inside "location", "song", ...), which had effectively disabled the inverted-boolean filter. Signed-off-by: Joe Khawand <93840910+Joe-Khawand@users.noreply.github.com> --- needle/dataset/dataset.py | 12 +- needle/dataset/generate.py | 30 +++- needle/model/architecture.py | 12 +- needle/model/run.py | 33 ++-- needle/training/eval.py | 319 +++++++++++++++++++++-------------- needle/training/optim.py | 39 +++-- needle/training/pretrain.py | 3 +- needle/training/train.py | 3 +- 8 files changed, 282 insertions(+), 169 deletions(-) diff --git a/needle/dataset/dataset.py b/needle/dataset/dataset.py index 3d4a1ea..d1a01bd 100644 --- a/needle/dataset/dataset.py +++ b/needle/dataset/dataset.py @@ -71,7 +71,12 @@ def download_hf_split(split="train", repo_id=HF_DATASET_REPO): def _mark_json_value(s, char_w, key, value_str, weight): - """Find '"key": "value_str"' or '"key": value_str' in s, mark value chars.""" + """Mark value chars for every '"key": "value_str"' / '"key": value_str' in s. + + Marks all occurrences, not just the first, so repeated names/values across + parallel calls are weighted rather than left at base weight. + """ + found = False pattern_str = f'"{_re.escape(key)}"\\s*:\\s*"{_re.escape(value_str)}"' for m in _re.finditer(pattern_str, s): tail = s[m.start() + len(f'"{key}"'):m.end()] @@ -79,8 +84,10 @@ def _mark_json_value(s, char_w, key, value_str, weight): val_start = m.start() + len(f'"{key}"') + val_offset val_end = val_start + len(value_str) char_w[val_start:val_end] = np.maximum(char_w[val_start:val_end], weight) + found = True + if found: return - + pattern_ns = f'"{_re.escape(key)}"\\s*:\\s*{_re.escape(value_str)}' for m in _re.finditer(pattern_ns, s): colon_offset = s[m.start():m.end()].index(':') @@ -89,7 +96,6 @@ def _mark_json_value(s, char_w, key, value_str, weight): val_start += 1 val_end = m.end() char_w[val_start:val_end] = np.maximum(char_w[val_start:val_end], weight) - return def _mark_json_key_in_args(s, char_w, key, weight): diff --git a/needle/dataset/generate.py b/needle/dataset/generate.py index cc358b0..a055359 100644 --- a/needle/dataset/generate.py +++ b/needle/dataset/generate.py @@ -1247,6 +1247,27 @@ def _grounding_check(pname, pval, pdesc, query, query_lower): return True +def _boolean_enable_conflict(query_lower, pval): + """True if a boolean 'enabled' value clearly contradicts the query intent. + + Matches whole words/phrases, not substrings: "on" in "location"/"song" + otherwise makes the on-intent always fire and disables the filter. + """ + words = set(re.findall(r"[a-z']+", query_lower)) + off_words = {"off", "disable", "disabled", "stop", "deactivate", "without"} + on_words = {"on", "enable", "enabled", "start", "activate"} + off_phrases = ("turn off", "switch off", "shut off", "turn it off", "don't") + on_phrases = ("turn on", "switch on", "turn it on") + wants_off = bool(words & off_words) or any(p in query_lower for p in off_phrases) + wants_on = bool(words & on_words) or any(p in query_lower for p in on_phrases) + # Only flag clear contradictions — if ambiguous, allow. + if wants_off and not wants_on and pval is True: + return True + if wants_on and not wants_off and pval is False: + return True + return False + + def _semantic_check(tool_name, args, schema, query, call_type="single"): """Lightweight rule-based semantic validation of argument values. @@ -1297,14 +1318,7 @@ def _semantic_check(tool_name, args, schema, query, call_type="single"): # check alignment with query intent if expected_type == "boolean" and isinstance(pval, bool): if pname == "enabled": - disable_words = ("off", "disable", "stop", "don't", "no ", "without") - enable_words = ("on", "enable", "start", "turn on", "activate") - query_wants_off = any(w in query_lower for w in disable_words) - query_wants_on = any(w in query_lower for w in enable_words) - # Only reject clear contradictions — if ambiguous, allow - if query_wants_off and not query_wants_on and pval is True: - return False - if query_wants_on and not query_wants_off and pval is False: + if _boolean_enable_conflict(query_lower, pval): return False # Grounding check: argument values must be traceable to the query. diff --git a/needle/model/architecture.py b/needle/model/architecture.py index c62cc4d..390716f 100644 --- a/needle/model/architecture.py +++ b/needle/model/architecture.py @@ -355,11 +355,19 @@ def encode(self, src, src_mask=None): """Backward-compatible alias for encode_text.""" return self.encode_text(src, src_mask=src_mask) - def decode(self, tgt, encoder_out, self_mask=None, cross_mask=None, deterministic=True): - """Decode from encoder output with cross_mask for variable-length encoder output.""" + def decode(self, tgt, encoder_out, self_mask=None, cross_mask=None, deterministic=True, cur_pos=None): + """Decode from encoder output with cross_mask for variable-length encoder output. + + When cur_pos (a possibly-traced scalar) is given, only that position is + projected to the vocabulary — returns (B, 1, V) and skips the full-buffer + (B, T, d) @ (d, V) matmul that dominates each decode step. cur_pos=None + keeps the training path unchanged. + """ x = self.embedding(tgt) * self.embed_scale rope = self._rope(tgt.shape[1]) x = self.decoder(x, encoder_out, self_mask=self_mask, cross_mask=cross_mask, rope=rope, deterministic=deterministic) + if cur_pos is not None: + x = jax.lax.dynamic_slice_in_dim(x, cur_pos, 1, axis=1) logits = x.astype(jnp.float32) @ self.embedding.embedding.T return logits diff --git a/needle/model/run.py b/needle/model/run.py index 954022e..64d1b93 100644 --- a/needle/model/run.py +++ b/needle/model/run.py @@ -71,10 +71,10 @@ def _get_decode_fn(model, max_gen_len): tgt_mask = make_causal_mask(max_gen_len) @jax.jit - def decode_step(params, dec_buffer, encoder_out, cross_mask): + def decode_step(params, dec_buffer, encoder_out, cross_mask, cur_pos): return model.apply( {"params": params}, dec_buffer, encoder_out, - self_mask=tgt_mask, cross_mask=cross_mask, method="decode", + self_mask=tgt_mask, cross_mask=cross_mask, cur_pos=cur_pos, method="decode", ) _decode_fn_cache[key] = decode_step @@ -140,10 +140,11 @@ def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MA sys.stdout.write(f"\n") sys.stdout.flush() - logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) + streamed = "" + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(0, dtype=jnp.int32)) for i in range(0, max_gen_len - 1): - next_logits = logits[0, i] + next_logits = logits[0, 0] if constrained_decoder and constrained_decoder.is_active(0): logits_np = np.array(next_logits) @@ -162,10 +163,18 @@ def generate(model, params, tokenizer, query, tools="[]", max_gen_len=DEFAULT_MA dec_buffer = dec_buffer.at[0, i + 1].set(next_token) if stream: - sys.stdout.write(tokenizer.decode([next_token])) - sys.stdout.flush() - - logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) + # Emit the decode delta of the whole sequence, not the token alone: + # per-token decode drops SentencePiece spaces and splits multibyte + # pieces. Hold until the trailing char completes (no U+FFFD). + full = tokenizer.decode(generated_tokens) + if not full.endswith("�"): + delta = full[len(streamed):] + if delta: + sys.stdout.write(delta) + sys.stdout.flush() + streamed = full + + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(i + 1, dtype=jnp.int32)) if stream: sys.stdout.write("\n") @@ -229,18 +238,18 @@ def generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=DE from .constrained import build_constrained_decoder constrained_decoder = build_constrained_decoder(tools_list, tokenizer) - logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(0, dtype=jnp.int32)) for pos in range(0, max_gen_len - 1): for i in range(B): if finished[i]: continue if constrained_decoder and constrained_decoder.is_active(i): - logits_np = np.array(logits[i, pos]) + logits_np = np.array(logits[i, 0]) logits_np = constrained_decoder.constrain_logits(logits_np, i) next_token = int(np.argmax(logits_np)) else: - next_token = int(jnp.argmax(logits[i, pos])) + next_token = int(jnp.argmax(logits[i, 0])) if constrained_decoder: constrained_decoder.update(i, next_token) if next_token == eos_id: @@ -252,7 +261,7 @@ def generate_batch(model, params, tokenizer, queries, tools_list, max_gen_len=DE if all(finished): break - logits = decode_fn(params, dec_buffer, encoder_out, enc_mask) + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(pos + 1, dtype=jnp.int32)) results = [] for i in range(B): diff --git a/needle/training/eval.py b/needle/training/eval.py index 1b0d1dc..69f5d44 100644 --- a/needle/training/eval.py +++ b/needle/training/eval.py @@ -1,13 +1,15 @@ import argparse +import json import math import time +from collections import Counter import jax import jax.numpy as jnp import numpy as np import optax -from ..dataset.dataset import get_tokenizer, load_tool_calls, load_prepared_data, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_DEC_LEN, DEFAULT_MAX_GEN_LEN +from ..dataset.dataset import get_tokenizer, load_tool_calls, load_prepared_data, to_snake_case, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_DEC_LEN, DEFAULT_MAX_GEN_LEN from ..model.architecture import ( SimpleAttentionNetwork, TransformerConfig, @@ -72,7 +74,7 @@ def measure_throughput(model, params, tokenizer, num_runs=10, prompt='What is th ) 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) + decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(0, dtype=jnp.int32)) tokens_generated = [] latencies = [] @@ -86,19 +88,19 @@ def measure_throughput(model, params, tokenizer, num_runs=10, prompt='What is th 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) + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(0, dtype=jnp.int32)) 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() + next_token = jax.random.categorical(sample_rng, logits[0, 0]).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) + logits = decode_fn(params, dec_buffer, encoder_out, enc_mask, jnp.array(i + 1, dtype=jnp.int32)) elapsed = time.perf_counter() - start tokens_generated.append(num_tokens) @@ -178,74 +180,86 @@ def compute_wer(hypotheses, references): 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 ..model.run import generate_batch, normalize_tools, restore_tool_names - from ..dataset.dataset 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) - +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) + v = str(f) + except ValueError: + pass + s = v.strip().lower() + if s.startswith("at "): + s = s[3:].strip() + 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): + """Canonical (name, normalized-args) key for a single call, or None.""" + 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) + + +def classify_bfcl(ref_calls, n_tools): + """Bucket an example into a BFCL-style category from its reference calls. + + - irrelevance : no call expected + - simple : one call, one candidate tool + - multiple : one call, several candidate tools (selection) + - parallel : several calls, all to the same function + - parallel_multiple: several calls spanning >=2 distinct functions + """ + calls = [c for c in ref_calls if isinstance(c, dict)] + n_calls = len(calls) + if n_calls == 0: + return "irrelevance" + n_unique = len({c.get("name") for c in calls}) + if n_calls == 1: + return "simple" if n_tools <= 1 else "multiple" + return "parallel" if n_unique == 1 else "parallel_multiple" + + +def score_tool_calls(ds, all_preds): + """Score decoded predictions against references. Pure function (no model). + + ``ds`` is an iterable of {query, tools, answers}; ``all_preds`` is the + parallel list of decoded prediction strings. Returns the metrics dict. + Split out from generation so the scoring logic can be unit-tested on + hand-built parallel / irrelevance cases without loading a checkpoint. + """ total = 0 exact_match = 0 - tp_names = 0 - fp_names = 0 - fn_names = 0 - tp_calls = 0 - fp_calls = 0 - fn_calls = 0 + tp_names = fp_names = fn_names = 0 + tp_calls = fp_calls = 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 + empty_ref = empty_pred = 0 + args_correct = args_total = 0 + halluc_params = total_pred_params = 0 + missing_params = total_ref_params = 0 + correct_values = matched_params = 0 samples = [] - failures = [] # (query, tools, ref, pred, reasons) + failures = [] - 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) + CATS = ("simple", "multiple", "parallel", "parallel_multiple") + cat_total = {c: 0 for c in CATS} + cat_correct = {c: 0 for c in CATS} + irrel_total = irrel_correct = 0 # ref empty; correct == abstained + rel_total = rel_correct = 0 # ref non-empty; correct == called something for i, (ex, pred_text) in enumerate(zip(ds, all_preds)): ref_text = ex["answers"] @@ -255,14 +269,13 @@ def call_key(c): ref_calls = json.loads(ref_text) except (json.JSONDecodeError, TypeError): ref_calls = [] - - # Normalize ref tool names to snake_case for consistent comparison + if not isinstance(ref_calls, list): + ref_calls = [] 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) + # Normalize pred tool names to snake_case for comparison. try: pred_calls_raw = json.loads(pred_text) if isinstance(pred_calls_raw, list): @@ -275,15 +288,18 @@ def call_key(c): except (json.JSONDecodeError, TypeError): pass - # Also normalize tool definitions for param validation try: tool_defs_raw = json.loads(ex["tools"]) + if not isinstance(tool_defs_raw, list): + tool_defs_raw = [] 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): + tool_defs_raw = [] tools_normalized = ex["tools"] + n_tools = len(tool_defs_raw) try: pred_calls = json.loads(pred_text) @@ -304,31 +320,56 @@ def call_key(c): if pred_is_empty: empty_pred += 1 + # Exact (order-independent, multiset-correct) match. if ref_is_empty and pred_is_empty: - exact_match += 1 + is_exact = True 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_sorted = sorted(k for k in (_call_key(c) for c in ref_calls) if k) + pred_sorted = sorted(k for k in (_call_key(c) for c in pred_calls) if k) + is_exact = (ref_sorted == pred_sorted + and len(ref_sorted) == len(ref_calls) + and len(pred_sorted) == len(pred_calls)) + else: + is_exact = False + if is_exact: + exact_match += 1 + # Relevance / irrelevance + per-category AST accuracy. + if ref_is_empty: + irrel_total += 1 + if pred_is_empty: + irrel_correct += 1 + else: + rel_total += 1 + if not pred_is_empty: + rel_correct += 1 + cat = classify_bfcl(ref_calls, n_tools) + if cat in cat_total: + cat_total[cat] += 1 + if is_exact: + cat_correct[cat] += 1 + + # Name-level P/R (set semantics fine for names). 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) + # Call-level P/R as MULTISETS — identical parallel calls must not collapse. + ref_counter = Counter(k for k in (_call_key(c) for c in ref_calls) if k) + pred_counter = Counter(k for k in (_call_key(c) for c in pred_calls) if k) + tp_calls += sum((ref_counter & pred_counter).values()) + fp_calls += sum((pred_counter - ref_counter).values()) + fn_calls += sum((ref_counter - pred_counter).values()) + # Consumable ref-arg instances per name. 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", {})) + ref_by_name.setdefault(c["name"], []).append(c.get("arguments", {}) or {}) + + # args_acc: each pred vs ANY ref instance of the same name. for c in pred_calls: if isinstance(c, dict) and "name" in c and c["name"] in ref_by_name: args_total += 1 @@ -336,83 +377,77 @@ def call_key(c): if any(pa == json.dumps(_normalize_args(ra), sort_keys=True) for ra in ref_by_name[c["name"]]): args_correct += 1 + # Param-level metrics via greedy bipartite matching to DISTINCT ref instances. 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} + 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 = {} + + used = {name: [False] * len(lst) for name, lst in ref_by_name.items()} + matched_pairs = [] # (pred_call, matched_ref_args) — for failure diagnosis 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) + p_args = c.get("arguments", {}) or {} + p_keys = set(p_args.keys()) + if cname in tool_param_map: + total_pred_params += len(p_keys) + halluc_params += len(p_keys - tool_param_map[cname]) 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()) + cand = [j for j in range(len(ref_by_name[cname])) if not used[cname][j]] + if not cand: + continue + pnorm = json.dumps(_normalize_args(p_args), sort_keys=True) + + def _match_score(j, _cname=cname, _pnorm=pnorm, _pk=p_keys): + ra = ref_by_name[_cname][j] + exact = json.dumps(_normalize_args(ra), sort_keys=True) == _pnorm + return (exact, len(_pk & set(ra.keys()))) + + j = max(cand, key=_match_score) + used[cname][j] = True + ref_args = ref_by_name[cname][j] + matched_pairs.append((c, ref_args)) + r_keys = set(ref_args.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): + if (json.dumps(_normalize_value(p_args[k]), sort_keys=True) + == json.dumps(_normalize_value(ref_args[k]), 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) - ) + # Failure diagnosis. 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") - + if not pred_calls and pred_text.strip() not in ("", "[]"): + 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 + for c, ref_args in matched_pairs: cname = c["name"] - if cname not in ref_by_name: - continue - pred_args = c.get("arguments", {}) - ref_args = ref_by_name[cname][0] + pred_args = c.get("arguments", {}) or {} 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): + if (json.dumps(_normalize_value(pv), sort_keys=True) + != json.dumps(_normalize_value(rv), 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], @@ -444,11 +479,32 @@ def call_key(c): "call_f1": call_f1, "empty_ref_pct": empty_ref / max(total, 1), "empty_pred_pct": empty_pred / max(total, 1), + # BFCL-style AST accuracy per category (exact match within each bucket). + "category_counts": dict(cat_total), + "category_accuracy": {c: (cat_correct[c] / cat_total[c]) if cat_total[c] else None for c in CATS}, + "irrelevance_accuracy": (irrel_correct / irrel_total) if irrel_total else None, + "false_trigger_rate": ((irrel_total - irrel_correct) / irrel_total) if irrel_total else None, + "relevance_accuracy": (rel_correct / rel_total) if rel_total else None, + "irrelevance_n": irrel_total, + "relevance_n": rel_total, "samples": samples, "failures": failures, } +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.""" + from ..model.run import generate_batch + + 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) + return score_tool_calls(ds, all_preds) + + 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. @@ -620,6 +676,17 @@ def _classify(ex): 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%}") + # BFCL-style per-category AST accuracy + abstention/relevance. + cat_acc = tc.get("category_accuracy", {}) + cat_n = tc.get("category_counts", {}) + for c in ("simple", "multiple", "parallel", "parallel_multiple"): + if cat_acc.get(c) is not None: + print(f" {('AST ' + c):<16} {cat_acc[c]:>10.1%} (n={cat_n.get(c, 0)})") + if tc.get("relevance_accuracy") is not None: + print(f" Relevance acc {tc['relevance_accuracy']:>10.1%} (n={tc.get('relevance_n', 0)})") + if tc.get("irrelevance_accuracy") is not None: + print(f" Irrelevance acc {tc['irrelevance_accuracy']:>10.1%} (n={tc.get('irrelevance_n', 0)})") + print(f" False trigger {tc['false_trigger_rate']:>10.1%}") if tc["samples"]: print(f" samples:") for query, ref, pred in tc["samples"][:3]: diff --git a/needle/training/optim.py b/needle/training/optim.py index be8890b..fdf3b49 100644 --- a/needle/training/optim.py +++ b/needle/training/optim.py @@ -8,7 +8,11 @@ def _newton_schulz(G, steps=5): - """Approximate polar decomposition via Newton-Schulz iteration.""" + """Approximate polar decomposition via Newton-Schulz, with aspect-ratio scaling. + + Flax kernels are (fan_in, fan_out); the max(1, fan_out/fan_in)**0.5 factor + keeps the update RMS consistent across non-square weights. + """ a, b, c = 3.4445, -4.7750, 2.0315 orig_dtype = G.dtype G = G.astype(jnp.float32) @@ -22,7 +26,8 @@ def _newton_schulz(G, steps=5): X = a * X + B @ X if transposed: X = X.T - return X.astype(orig_dtype) + scale = jnp.sqrt(jnp.maximum(1.0, G.shape[1] / G.shape[0])) + return (X * scale).astype(orig_dtype) class MuonState(NamedTuple): @@ -30,26 +35,28 @@ class MuonState(NamedTuple): def scale_by_muon(momentum=0.95, ns_steps=5): - """Muon gradient transform: orthogonalize 2D+ grads, then Nesterov momentum.""" + """Muon gradient transform: Nesterov momentum on the raw grad, then orthogonalize. + + Orthogonalizing before momentum accumulates a sum of orthogonal matrices, + which is not itself orthogonal — so the buffer holds raw grads and only the + Nesterov-blended update is passed through Newton-Schulz. + """ def init_fn(params): return MuonState(mu=jax.tree.map(jnp.zeros_like, params)) + def ortho(g): + if g.ndim == 3: + return jax.vmap(lambda m: _newton_schulz(m, steps=ns_steps), in_axes=(0,))(g) + if g.ndim == 2: + return _newton_schulz(g, steps=ns_steps) + return g + def update_fn(updates, state, params=None): del params - - def ortho(g): - if g.ndim == 3: - return jax.vmap(_newton_schulz, in_axes=(0,))(g) - if g.ndim == 2: - return _newton_schulz(g, steps=ns_steps) - return g - - ortho_g = jax.tree.map(ortho, updates) - new_mu = jax.tree.map(lambda m, g: momentum * m + g, state.mu, ortho_g) - new_updates = jax.tree.map( - lambda g, m: g + momentum * m, ortho_g, new_mu - ) + new_mu = jax.tree.map(lambda m, g: momentum * m + g, state.mu, updates) + blended = jax.tree.map(lambda g, m: g + momentum * m, updates, new_mu) + new_updates = jax.tree.map(ortho, blended) return new_updates, MuonState(mu=new_mu) return optax.GradientTransformation(init_fn, update_fn) diff --git a/needle/training/pretrain.py b/needle/training/pretrain.py index 6c1bf83..df030a7 100644 --- a/needle/training/pretrain.py +++ b/needle/training/pretrain.py @@ -156,7 +156,8 @@ def loss_fn(params): ce = jnp.sum( optax.softmax_cross_entropy_with_integer_labels(logits_f32, tgt_out) * padding_mask ) / num_tokens - z_loss = 1e-4 * jnp.mean(jax.nn.logsumexp(logits_f32, axis=-1) ** 2) + # Mask z-loss to real tokens (like CE), else it scales with padding fraction. + z_loss = 1e-4 * jnp.sum((jax.nn.logsumexp(logits_f32, axis=-1) ** 2) * padding_mask) / num_tokens return ce + z_loss loss, grads = jax.value_and_grad(loss_fn)(state.params) diff --git a/needle/training/train.py b/needle/training/train.py index 657559f..21d2319 100644 --- a/needle/training/train.py +++ b/needle/training/train.py @@ -78,7 +78,8 @@ def _text_loss_fn(state, params, src, tgt_in, tgt_out, rng, loss_mask, do_quanti ce_loss = jnp.sum( optax.softmax_cross_entropy_with_integer_labels(logits_f32, tgt_out) * mask ) / num_tokens - z_loss = 1e-4 * jnp.mean(jax.nn.logsumexp(logits_f32, axis=-1) ** 2) + # Mask z-loss to real tokens (like CE), else it scales with padding fraction. + z_loss = 1e-4 * jnp.sum((jax.nn.logsumexp(logits_f32, axis=-1) ** 2) * padding_mask) / num_tokens return ce_loss + z_loss