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>

Author: Joe-Khawand

needle/dataset/dataset.py

@@ -71,16 +71,23 @@ 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()]
         val_offset = tail.index(f'"{value_str}"') + 1
         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):

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.

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
 

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):