ActionMan-0.1 research preview

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/arena-mining.jpg filter=lfs diff=lfs merge=lfs -text
+assets/how-it-works.jpg filter=lfs diff=lfs merge=lfs -text
+assets/latency_vs_output_size.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,15 @@
+ActionMan-0.1 weights: research preview licence
+
+Copyright (c) 2026 Shreyash Kumar Singh
+
+The model weights in this repository (model.safetensors) are released for research and evaluation use only.
+
+They were trained on data that includes Salesforce's xlam-function-calling-60k (CC-BY-4.0), whose dataset card states
+"This release is for research purposes only." Commercial use of these weights is not permitted under this preview licence.
+
+Other training data: Team-ACE/ToolACE (Apache-2.0) and NousResearch/hermes-function-calling-v1 (Apache-2.0).
+Base model: answerdotai/ModernBERT-large (Apache-2.0).
+
+The inference code in code/actionman/ is released under the MIT licence (see the GitHub repository).
+
+THE WEIGHTS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND.

README.md

@@ -0,0 +1,148 @@
+---
+license: other
+license_name: actionman-research-preview
+license_link: LICENSE
+library_name: pytorch
+base_model: answerdotai/ModernBERT-large
+tags:
+- function-calling
+- tool-calling
+- tool-use
+- non-autoregressive
+- masked-diffusion
+- modernbert
+- research
+datasets:
+- Salesforce/xlam-function-calling-60k
+- Team-ACE/ToolACE
+- NousResearch/hermes-function-calling-v1
+---
+
+# ActionMan-0.1
+
+**A 0.4B tool caller that doesn't write JSON.**
+- One encoder read picks the tool.
+- Masked diffusion fills every argument at once, in two or three passes.
+
+ActionMan started as a question: decision models answer in milliseconds but can only choose, while LLMs can act but write one token
+at a time. Can a decision-style encoder also write the arguments?
+
+> **Early research preview.**
+> - **Speed:** on one RTX 4090 at batch 1 it answers in **17 ms** (p50), where the same-data Qwen3-0.6B served by tuned vLLM takes
+>   **121 ms**. That is 7x for one call and 16x for three or more.
+> - **Accuracy:** it is **less accurate** than that LLM (BFCL v4 Non-Live 66.0 vs 75.7).
+> - **Scope:** it is single-turn only, and it is slow on a laptop CPU.
+>
+> It is fast and well calibrated; it is not a drop-in replacement for an LLM tool caller.
+
+![how it works](assets/how-it-works.jpg)
+
+## How it works
+
+- **Encoder:** ModernBERT-large (395M) reads the tools (as compact signatures), the query, and up to four `[CALL]` markers.
+- **Decision head:** each marker picks a tool or none, by dot product against the tool-name tokens. A length head sizes each argument.
+- **Slot block:** it is appended after the input and filled in parallel by mask-predict. Slots below 0.85 confidence are re-masked;
+  most requests finish in 2-3 passes. The input is encoded once and cached; the passes run as CUDA graphs.
+- **Copy, pick, list:**
+  - A copy switch points at a span of the query instead of spelling the value.
+  - Booleans and enums (up to 16 options) are picked from their options.
+  - Lists get a count head and item slots.
+
+About 410M unique parameters. The decoder is tied to the embeddings, but it is stored twice in the bf16 file.
+
+## Results
+
+| BFCL v4, official harness (`bfcl-eval==2026.3.23`) | size | Non-Live | Live | Irrelevance |
+|---|---|---|---|---|
+| **ActionMan-0.1** | 0.4B | 66.0 | 45.1 | **81.2** |
+| Hammer2.1-0.5b | 0.49B | 65.9 | 54.7 | 80.7 |
+| granite-4.0-350m | 0.35B | 67.7 | 46.7 | 61.1 |
+| Qwen3-0.6B (FC) | 0.75B | 72.9 | 58.8 | 80.9 |
+| Qwen3-0.6B fine-tuned on ActionMan's data | 0.75B | 75.7 | 54.1 | 79.1 |
+| Qwen3-1.7B (FC) | 2.0B | 82.7 | 74.7 | 76.8 |
+
+| Latency, batch 1, RTX 4090, 500 BFCL v4 prompts, end to end | p50 | p95 | correct on the 500 |
+|---|---|---|---|
+| **ActionMan-0.1, compiled** | **17.0 ms** | 28.9 ms | 59.2% |
+| Qwen3-0.6B fine-tuned, vLLM 0.11 (tuned) | 121 ms | 347 ms | 68.6% |
+| Qwen3-0.6B fine-tuned, llama.cpp Q8_0, GPU | 159 ms | 336 ms | 69.0% |
+| 8 server CPU threads: ActionMan fp32 vs llama.cpp Q8_0 | 0.45 s | | 59.0% vs 66.5% |
+
+![latency](assets/latency_vs_output_size.png)
+
+- **The harness** reproduces the public leaderboard for Hammer and granite within 0.25 Non-Live points.
+- **No BFCL Overall:** multi-turn and agentic categories were not run.
+- **The full report** has per-category scores with confidence intervals, error classes, and our own 2,340-prompt eval:
+  [docs/results/benchmark.md](https://github.com/shing1Sks/ActionMan/blob/dev/docs/results/benchmark.md).
+
+## Usage
+
+```bash
+pip install "git+https://github.com/shing1Sks/ActionMan"        # torch, transformers>=5, safetensors, huggingface_hub
+```
+
+```python
+from actionman.api import ActionMan
+
+am = ActionMan.from_pretrained("shing-dev/ActionMan-0.1")    # compiles on GPU (one-off warm-up)
+tools = [{"type": "function", "function": {
+    "name": "get_weather", "description": "Current weather for a city",
+    "parameters": {"type": "object", "properties": {
+        "city": {"type": "string"}, "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}},
+        "required": ["city"]}}}]
+am.predict("What's the weather in Paris in celsius?", tools)
+# [{"name": "get_weather", "arguments": {"city": "Paris", "unit": "celsius"}, "confidence": 0.9812}]
+am.predict("Tell me a joke", tools)
+# []
+```
+
+- **Without pip:** the inference code is also in this repo under `code/actionman/`. Add `code/` to `sys.path`.
+- **No call:** an empty list means no tool fits. `require=True` forces a call, like `tool_choice="required"`.
+- **`confidence`** is the minimum over the tool and every argument.
+
+## Calibration: knowing when it doesn't know
+
+![reliability](assets/reliability_bfcl_v4.png)
+
+| | our eval (2,340 prompts) | BFCL v4 Python |
+|---|---|---|
+| accuracy, answering everything | 78.6 | 56.1 |
+| accuracy when the least confident 10% abstain | 85.1 | 62.1 |
+| accuracy when the least confident 20% abstain | 89.5 | 66.9 |
+
+Confidence ranks errors well, but it is over-confident on BFCL (ECE 0.28, vs 0.11 on our eval). Use it as a threshold, not as a probability.
+
+## Limits
+
+- **Single turn, English.** At most 4 calls per request, lists of up to 7 items, and an input of about 896 tokens.
+- **No reasoning.**
+  - Values that are not in the text are its main error: "until 11 PM" -> `23`, unit conversion, world knowledge.
+  - On long, real user prompts (BFCL Live) it falls well behind small LLMs.
+- **CPU:** 0.45 s per call on 8 server cores, and several seconds on a laptop CPU. The speed claim is a GPU claim.
+- **The arena:** ActionMan executes instructions; it does not plan. Given a description of a game state instead of an instruction,
+  it loops (1 of 12 mining races won).
+
+## Training
+
+- **Data:** 5 epochs over 80,459 rows:
+  - xLAM-60k, 58,000;
+  - ToolACE, 10,359;
+  - Hermes v1, 1,100;
+  - no-call negatives, 11,000 (built from xLAM queries).
+- **Hardware and cost:** one RTX 4090, about 2 hours. The whole research, every run included, cost about $16 of rented GPU time,
+  plus $2 for the benchmarks.
+
+## Licence
+
+- **The weights are a research preview, for research use only.** The training data includes Salesforce's xLAM-60k, whose card
+  says "This release is for research purposes only."
+- **ToolACE and Hermes are Apache-2.0.**
+- **The inference code is MIT.**
+- **Open weights** would need a retrain without xLAM, which is planned.
+
+## Links
+
+- **Code, benchmarks, arena:** [github.com/shing1Sks/ActionMan](https://github.com/shing1Sks/ActionMan)
+- **The research film:** [actionman-research-film.mp4](https://github.com/shing1Sks/ActionMan/releases/download/v0.1/actionman-research-film.mp4)
+
+Research by Shreyash Kumar Singh with an AI coding agent (Claude).

code/actionman/__init__.py

@@ -0,0 +1,11 @@
+"""ActionMan: decision encoder + masked-diffusion argument filler."""
+import os
+import time
+
+
+def beat():
+    """Touch $HEARTBEAT (if set) so the box watchdog knows we're alive."""
+    p = os.environ.get("HEARTBEAT")
+    if p:
+        with open(p, "w") as f:
+            f.write(str(time.time()))

code/actionman/api.py

@@ -0,0 +1,53 @@
+"""Public API: load a released ActionMan and turn (query, tools) into tool calls.
+
+    from actionman.api import ActionMan
+    am = ActionMan.from_pretrained("shing-dev/ActionMan-0.1")          # or a local export directory
+    am.predict("Weather in Paris in celsius?", tools)                   # tools: OpenAI-style function schemas
+    # -> [{"name": "get_weather", "arguments": {"city": "Paris", "unit": "celsius"}, "confidence": 0.97}]
+"""
+from __future__ import annotations
+
+import json
+import os
+
+import torch
+
+from . import cached as C
+from .features import load_tokenizer
+from .infer import predict as _predict
+from .model import ActionMan as _Model
+from .schema import Example, from_openai_tool
+
+
+class ActionMan:
+    def __init__(self, model, tok, flags: dict, device: str):
+        self.model, self.tok, self.flags, self.device = model, tok, flags, device
+
+    @classmethod
+    def from_pretrained(cls, path: str, device: str | None = None, compile: bool | None = None, token: str | None = None):
+        """path: a local export dir or a Hugging Face repo id. compile (default: on GPU): padded static encoders
+        under torch.compile, about 4x faster per call after a one-off warm-up."""
+        from safetensors.torch import load_file
+        from transformers import ModernBertConfig
+        if not os.path.isdir(path):
+            from huggingface_hub import snapshot_download
+            path = snapshot_download(path, token=token)
+        cfg = json.load(open(os.path.join(path, "config.json")))
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        tok = load_tokenizer(path)
+        model = _Model.from_encoder(ModernBertConfig.from_dict(cfg["encoder"]), tok, v3=cfg["flags"]["v3"])
+        model.load_state_dict({k: v.float() for k, v in load_file(os.path.join(path, "model.safetensors")).items()})
+        model = model.to(device).eval()
+        if compile if compile is not None else device == "cuda":     # process-wide swap, as in evaluate --compile
+            C.encode_context, C.encode_slots = C.make_static(model.mlm.model)
+        return cls(model, tok, cfg["flags"], device)
+
+    @torch.no_grad()
+    def predict(self, query: str, tools: list[dict], T: int = 3, require: bool = False) -> list[dict]:
+        """Tool calls for `query`, best first. Empty list = no tool fits (unless require=True, like tool_choice="required").
+        confidence = min(tool confidence, every argument's confidence). It ranks errors well enough to abstain on
+        (docs/results/benchmark.md: the least confident 10% abstaining takes our eval from 78.6 to 85.1)."""
+        ex = Example(query, [from_openai_tool(t) for t in tools], [])
+        pred = _predict(self.model, self.tok, ex, T=T, device=self.device, require=require, **self.flags)
+        return [{"name": c.name, "arguments": c.args,
+                 "confidence": round(min([c.conf_tool, *c.conf_args.values()]), 4)} for c in pred.calls]

code/actionman/audit.py

@@ -0,0 +1,129 @@
+"""Data audit (v3 S0): how much argument mass can be copied from the query, picked from options, or is a list."""
+from __future__ import annotations
+
+import os
+from collections import Counter, defaultdict
+
+from .schema import param_options, value_to_text
+
+
+def query_offsets(tok, query: str) -> list[tuple[int, int]]:
+    """Char span per query token, with a leading space (BPE 'Ġ') trimmed off."""
+    offs = tok(query, add_special_tokens=False, return_offsets_mapping=True).offset_mapping
+    out = []
+    for a, b in offs:
+        while a < b and query[a].isspace():
+            a += 1
+        out.append((a, b))
+    return out
+
+
+def copy_span(tok, query: str, text: str, offs=None):
+    """First case-insensitive, token-aligned occurrence of `text` in `query` -> (start_tok, end_tok) or None."""
+    text = str(text).strip()
+    if not text:
+        return None
+    offs = offs if offs is not None else query_offsets(tok, query)
+    starts = {a: i for i, (a, b) in reversed(list(enumerate(offs))) if b > a}
+    ends = {b: i for i, (a, b) in enumerate(offs) if b > a}
+    low, t = query.lower(), text.lower()
+    at = low.find(t)
+    while at != -1:
+        s, e = starts.get(at), ends.get(at + len(t))
+        if s is not None and e is not None and e >= s:
+            return s, e
+        at = low.find(t, at + 1)
+    return None
+
+
+def _kind(v) -> str:
+    if isinstance(v, bool):
+        return "bool"
+    if isinstance(v, (int, float)):
+        return "num"
+    if isinstance(v, str):
+        return "str"
+    if isinstance(v, list):
+        return "list"
+    if isinstance(v, dict):
+        return "dict"
+    return "other"
+
+
+def audit(tok, examples) -> dict:
+    """Per-source counters over every gold argument value."""
+    stats: dict[str, Counter] = defaultdict(Counter)
+    for ex in examples:
+        c = stats[ex.source.split("_hardneg")[0].split("_neg")[0] if ex.source.startswith("xlam") else ex.source]
+        if not ex.calls:
+            continue
+        offs = query_offsets(tok, ex.query)
+        by_name = {t.name: t for t in ex.tools}
+        for call in ex.calls:
+            tool = by_name.get(call.name)
+            if tool is None:
+                continue
+            pmap = {p.name: p for p in tool.params}
+            for name, v in call.args.items():
+                p = pmap.get(name)
+                k = _kind(v)
+                c["args"] += 1
+                c[f"type_{k}"] += 1
+                n_tok = len(tok(value_to_text(v), add_special_tokens=False).input_ids)
+                if n_tok > 32:
+                    c["over_32_tokens"] += 1
+                opts = param_options(p) if p else None
+                if opts is not None:
+                    c["has_options"] += 1
+                    if any(value_to_text(o).lower() == value_to_text(v).lower() for o in opts):
+                        c["gold_in_options"] += 1
+                if k in ("str", "num"):
+                    c[f"{k}_total"] += 1
+                    if copy_span(tok, ex.query, value_to_text(v), offs) is not None:
+                        c[f"{k}_copyable"] += 1
+                if k == "list":
+                    c["list_items"] += len(v)
+                    if len(v) > 7:
+                        c["list_over_7"] += 1
+                    for it in v:
+                        c["list_item_total"] += 1
+                        if len(tok(value_to_text(it), add_special_tokens=False).input_ids) > 8:
+                            c["list_item_over_8_tokens"] += 1
+                        if isinstance(it, (str, int, float)) and not isinstance(it, bool) and \
+                                copy_span(tok, ex.query, value_to_text(it), offs) is not None:
+                            c["list_item_copyable"] += 1
+    return stats
+
+
+def _pct(a, b):
+    return f"{100 * a / b:.1f}%" if b else "-"
+
+
+def report(stats: dict) -> str:
+    rows = ["| source | args | str copyable | num copyable | bool/enum (options) | gold in options | list args | "
+            "list items copyable | item >8 tok | list >7 | value >32 tok |", "|" + "---|" * 11]
+    for src in sorted(stats):
+        c = stats[src]
+        if not c["args"]:
+            continue
+        rows.append(f"| {src} | {c['args']} | {_pct(c['str_copyable'], c['str_total'])} of {c['str_total']} | "
+                    f"{_pct(c['num_copyable'], c['num_total'])} of {c['num_total']} | "
+                    f"{_pct(c['has_options'], c['args'])} | {_pct(c['gold_in_options'], c['has_options'])} | "
+                    f"{_pct(c['type_list'], c['args'])} | {_pct(c['list_item_copyable'], c['list_item_total'])} | "
+                    f"{_pct(c['list_item_over_8_tokens'], c['list_item_total'])} | {c['list_over_7']} | "
+                    f"{_pct(c['over_32_tokens'], c['args'])} |")
+    return "\n".join(rows)
+
+
+if __name__ == "__main__":
+    from .data import load_jsonl
+    from .features import load_tokenizer
+    tok = load_tokenizer()
+    parts = []
+    for f in ("train_v2.jsonl", "test.jsonl", "bfcl.jsonl"):
+        st = audit(tok, load_jsonl(os.path.join("data/prepared", f)))
+        parts.append(f"## {f}\n\n{report(st)}\n")
+    out = "# v3 data audit (S0)\n\nShares are over gold argument values. 'copyable' = the value text occurs in the " \
+          "query, case-insensitively, on token boundaries.\n\n" + "\n".join(parts)
+    open("docs/results/audit.md", "w", encoding="utf-8").write(out)
+    print(out)

code/actionman/cached.py

@@ -0,0 +1,177 @@
+"""Block-causal 'cached' variant (Stage 4).
+
+Sequence = [pass-A context incl. [CALL]x4 [SEP]] [ARGS] slot block [SEP].
+Mask rule: context tokens see only context; slot-block tokens see everything.
+Consequences:
+  * training needs ONE forward (decision/length heads read context states that can't see the slots);
+  * at inference the context K/V of every layer is computed once (that IS pass A) and each refinement
+    pass runs only the ~10-40 slot tokens through the network.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb
+
+
+def block_masks(attn: torch.Tensor, ctx_len: torch.Tensor, window: int) -> dict:
+    """attn [B,L] (1 = real token), ctx_len [B]. Returns HF ModernBERT mask dict (bool, True = attend)."""
+    B, L = attn.shape
+    idx = torch.arange(L, device=attn.device)
+    valid_k = attn.bool()[:, None, None, :]
+    q_slot = (idx[None, :] >= ctx_len[:, None])[:, None, :, None]
+    k_ctx = (idx[None, :] < ctx_len[:, None])[:, None, None, :]
+    eye = torch.eye(L, dtype=torch.bool, device=attn.device)[None, None]
+    full = (valid_k & (q_slot | k_ctx)) | eye
+    win = ((idx[:, None] - idx[None, :]).abs() <= window)[None, None]
+    return {"full_attention": full, "sliding_attention": full & win}
+
+
+def _layer_qkv(layer, x, cos, sin):
+    attn = layer.attn
+    B, L, _ = x.shape
+    qkv = attn.Wqkv(layer.attn_norm(x)).view(B, L, 3, -1, attn.head_dim)
+    q, k, v = (t.transpose(1, 2) for t in qkv.unbind(dim=2))
+    q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
+    return q, k, v
+
+
+def _finish(layer, x, out):
+    B, L = x.shape[:2]
+    x = x + layer.attn.Wo(out.transpose(1, 2).reshape(B, L, -1))
+    return x + layer.mlp(layer.mlp_norm(x))
+
+
+@torch.no_grad()
+def encode_context(enc, ids: torch.Tensor):
+    """enc = ModernBertModel. ids [1,Lc]. Returns (final hidden [1,Lc,d], per-layer (k, v) cache)."""
+    cfg = enc.config
+    x = enc.embeddings(input_ids=ids)
+    L = ids.shape[1]
+    pos = torch.arange(L, device=ids.device)[None]
+    pe = {lt: enc.rotary_emb(x, pos, lt) for lt in set(cfg.layer_types)}
+    idx = torch.arange(L, device=ids.device)
+    win = ((idx[:, None] - idx[None, :]).abs() <= cfg.sliding_window)[None, None]
+    cache = []
+    for layer in enc.layers:
+        cos, sin = pe[layer.attention_type]
+        q, k, v = _layer_qkv(layer, x, cos, sin)
+        mask = win if layer.attention_type == "sliding_attention" else None
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5)
+        x = _finish(layer, x, out)
+        cache.append((k, v))
+    return enc.final_norm(x), cache
+
+
+@torch.no_grad()
+def encode_slots(enc, slot_ids: torch.Tensor, start: int, cache):
+    """Run only the slot-block tokens (absolute positions start..start+S-1) over the cached context K/V."""
+    cfg = enc.config
+    x = enc.embeddings(input_ids=slot_ids)
+    S = slot_ids.shape[1]
+    Lc = cache[0][0].shape[2]
+    qpos = torch.arange(start, start + S, device=slot_ids.device)
+    kpos = torch.cat([torch.arange(Lc, device=slot_ids.device), qpos])
+    pe = {lt: enc.rotary_emb(x, qpos[None], lt) for lt in set(cfg.layer_types)}
+    win = ((qpos[:, None] - kpos[None, :]).abs() <= cfg.sliding_window)[None, None]
+    for layer, (ck, cv) in zip(enc.layers, cache):
+        cos, sin = pe[layer.attention_type]
+        q, k, v = _layer_qkv(layer, x, cos, sin)
+        k, v = torch.cat([ck, k], dim=2), torch.cat([cv, v], dim=2)
+        mask = win if layer.attention_type == "sliding_attention" else None
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5)
+        x = _finish(layer, x, out)
+    return enc.final_norm(x)
+
+
+# ---------- static-shape variants for CUDA graphs / torch.compile (v3 S6) ----------
+# Lengths are padded up to a bucket and the real lengths are passed as tensors, so one compiled graph serves every
+# prompt in a bucket. Padded keys are masked out; outputs at padded positions are garbage and get sliced off.
+
+def _ctx_core(enc, ids, valid):
+    """ids [1,P] padded; valid = tensor([real length]). -> hidden [1,P,d], cache."""
+    cfg = enc.config
+    x = enc.embeddings(input_ids=ids)
+    P = ids.shape[1]
+    idx = torch.arange(P, device=ids.device)
+    pos = idx[None]
+    pe = {lt: enc.rotary_emb(x, pos, lt) for lt in set(cfg.layer_types)}
+    keyok = (idx < valid)[None, None, None, :]
+    eye = torch.eye(P, dtype=torch.bool, device=ids.device)[None, None]
+    full = keyok | eye
+    win = full & ((idx[:, None] - idx[None, :]).abs() <= cfg.sliding_window)[None, None]
+    cache = []
+    for layer in enc.layers:
+        cos, sin = pe[layer.attention_type]
+        q, k, v = _layer_qkv(layer, x, cos, sin)
+        mask = win if layer.attention_type == "sliding_attention" else full
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5)
+        x = _finish(layer, x, out)
+        cache.append((k, v))
+    return enc.final_norm(x), cache
+
+
+def _slot_core(enc, slot_ids, start, s_valid, c_valid, cache):
+    """slot_ids [1,Sp] padded; start/s_valid/c_valid tensors; cache from _ctx_core (padded context)."""
+    cfg = enc.config
+    x = enc.embeddings(input_ids=slot_ids)
+    Sp = slot_ids.shape[1]
+    Lc = cache[0][0].shape[2]
+    ar = torch.arange(Sp, device=slot_ids.device)
+    qpos = start + ar
+    kc = torch.arange(Lc, device=slot_ids.device)
+    kpos = torch.cat([kc, qpos])
+    keyok = torch.cat([kc < c_valid, ar < s_valid])[None, None, None, :]
+    pe = {lt: enc.rotary_emb(x, qpos[None], lt) for lt in set(cfg.layer_types)}
+    full = keyok | torch.cat([torch.zeros(Sp, Lc, dtype=torch.bool, device=slot_ids.device),
+                              torch.eye(Sp, dtype=torch.bool, device=slot_ids.device)], dim=1)[None, None]
+    win = full & ((qpos[:, None] - kpos[None, :]).abs() <= cfg.sliding_window)[None, None]
+    for layer, (ck, cv) in zip(enc.layers, cache):
+        cos, sin = pe[layer.attention_type]
+        q, k, v = _layer_qkv(layer, x, cos, sin)
+        k, v = torch.cat([ck, k], dim=2), torch.cat([cv, v], dim=2)
+        mask = win if layer.attention_type == "sliding_attention" else full
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5)
+        x = _finish(layer, x, out)
+    return enc.final_norm(x)
+
+
+def _bucket(n: int, b: int) -> int:
+    return -(-n // b) * b
+
+
+def make_static(enc, ctx_bucket: int = 256, slot_bucket: int = 128, compile: bool = True):
+    """-> (encode_context, encode_slots) drop-ins that pad to buckets (and run compiled, CUDA-graphed cores).
+    Buckets are coarse on purpose: every (context bucket, slot bucket) pair is a separate compiled graph."""
+    ctx_core, slot_core = _ctx_core, _slot_core
+    if compile:
+        import torch._dynamo as dynamo
+        for name in ("cache_size_limit", "recompile_limit"):
+            if hasattr(dynamo.config, name):
+                setattr(dynamo.config, name, 64)
+        ctx_core = torch.compile(_ctx_core, mode="reduce-overhead", dynamic=False, fullgraph=False)
+        slot_core = torch.compile(_slot_core, mode="reduce-overhead", dynamic=False, fullgraph=False)
+    pad_id = enc.config.pad_token_id or 0
+    state = {}
+
+    @torch.no_grad()
+    def encode_context_s(enc_, ids):
+        L = ids.shape[1]
+        P = _bucket(L, ctx_bucket)
+        ids_p = F.pad(ids, (0, P - L), value=pad_id)
+        h, cache = ctx_core(enc, ids_p, torch.tensor(L, device=ids.device))
+        h, cache = h.clone(), [(k.clone(), v.clone()) for k, v in cache]   # outlive the CUDA-graph buffers
+        state["c_valid"] = L
+        return h[:, :L], cache
+
+    @torch.no_grad()
+    def encode_slots_s(enc_, slot_ids, start, cache):
+        S = slot_ids.shape[1]
+        Sp = _bucket(S, slot_bucket)
+        ids_p = F.pad(slot_ids, (0, Sp - S), value=pad_id)
+        dev = slot_ids.device
+        out = slot_core(enc, ids_p, torch.tensor(start, device=dev), torch.tensor(S, device=dev),
+                        torch.tensor(state["c_valid"], device=dev), cache)
+        return out[:, :S].clone()
+
+    return encode_context_s, encode_slots_s

code/actionman/data.py

@@ -0,0 +1,220 @@
+"""Datasets, collation (random mask ratio = masked-diffusion training), splits, negatives."""
+from __future__ import annotations
+
+import json
+import random
+
+import torch
+
+from .features import build_pass_a, build_pass_b, gold_plan
+from .schema import Example, example_from_dict, example_to_dict
+
+
+def load_jsonl(path) -> list[Example]:
+    with open(path, encoding="utf-8") as f:
+        return [example_from_dict(json.loads(line)) for line in f if line.strip()]
+
+
+def save_jsonl(path, examples: list[Example]):
+    with open(path, "w", encoding="utf-8") as f:
+        for ex in examples:
+            f.write(json.dumps(example_to_dict(ex), ensure_ascii=False) + "\n")
+
+
+def featurize(tok, ex: Example, cached: bool = False, markers: bool = False, v3: bool = False) -> dict | None:
+    if v3:
+        from .v3feat import build_pass_a_v3, build_pass_b_v3, gold_plan_v3
+        a = build_pass_a_v3(tok, ex)
+        if a is None:
+            return None
+        b = None
+        if ex.calls and len(ex.calls) <= len(a["call_pos"]):
+            b = build_pass_b_v3(tok, ex, a, gold_plan_v3(tok, ex, a), with_targets=True, keep_context=cached,
+                                markers=markers)
+        return {"a": a, "b": b}
+    a = build_pass_a(tok, ex)
+    if a is None:
+        return None
+    b = None
+    if ex.calls and len(ex.calls) <= len(a["call_pos"]):
+        b = build_pass_b(tok, ex, a, gold_plan(a, ex), with_targets=True, keep_context=cached, markers=markers)
+    return {"a": a, "b": b}
+
+
+class FeatureDataset(torch.utils.data.Dataset):
+    def __init__(self, tok, examples: list[Example], cached: bool = False, markers: bool = False, v3: bool = False):
+        self.tok, self.examples, self.cached, self.markers, self.v3 = tok, examples, cached, markers, v3
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, i):
+        return featurize(self.tok, self.examples[i], self.cached, self.markers, self.v3)
+
+    def length(self, i) -> int:
+        """Cheap length proxy for bucketing (chars/3)."""
+        ex = self.examples[i]
+        n = len(ex.query) + sum(len(t.name) + len(t.description) + 20 * len(t.params) for t in ex.tools)
+        return n // 3 + 40
+
+
+def _pad(seqs, value):
+    L = max(len(s) for s in seqs)
+    ids = torch.full((len(seqs), L), value, dtype=torch.long)
+    attn = torch.zeros((len(seqs), L), dtype=torch.long)
+    for i, s in enumerate(seqs):
+        ids[i, : len(s)] = torch.tensor(s)
+        attn[i, : len(s)] = 1
+    return ids, attn
+
+
+def _mask_slots(b, rng, v3_mask: str = "clean"):
+    """Masked-diffusion corruption: random ratio r ~ U(0,1], mask that share of slot tokens, reveal the rest.
+    v3_mask="legacy" reproduces the original v3 run (s10): every anchor always masked, gate trained on all anchors."""
+    ids = list(b["input_ids"])
+    flat = [(p, t) for pos, tg in zip(b["slot_pos"], b["slot_targets"]) for p, t in zip(pos, tg)]
+    if not flat:                                   # v3: option-only slot block
+        return ids, []
+    n_mask = max(1, round(rng.uniform(0.0, 1.0) * len(flat)))
+    masked = set(rng.sample(range(len(flat)), n_mask))
+    always = set()
+    if b.get("anchors") and v3_mask == "legacy":
+        always = set(b["anchors"])
+        b.pop("_clean", None)
+    elif b.get("anchors"):
+        # v3: half the values are fully masked ("clean"); only their anchors train the copy gate / span pointer,
+        # matching inference, where the gate is read once with every slot masked. Other values follow the random
+        # mask, so revealed anchors (as in later refinement passes) are seen too.
+        clean = [pos for pos in b["slot_pos"] if rng.random() < 0.5]
+        always = {p for pos in clean for p in pos}
+        b["_clean"] = [pos[0] for pos in clean]
+    targets = []
+    for idx, (p, t) in enumerate(flat):
+        if idx in masked or p in always:
+            targets.append((p, t))
+        else:
+            ids[p] = t
+    return ids, targets
+
+
+def collate(items, tok, rng: random.Random | None = None, cached: bool = False, v3_mask: str = "clean"):
+    rng = rng or random
+    items = [it for it in items if it is not None]
+    if not items:
+        return None
+    if cached:
+        return _collate_cached(items, tok, rng, v3_mask)
+    A = [it["a"] for it in items]
+    a_ids, a_attn = _pad([a["input_ids"] for a in A], tok.pad_token_id)
+    J = max(max(len(a["t_pos"]) for a in A), 1)
+    t_pos = torch.full((len(A), J), -1, dtype=torch.long)
+    for i, a in enumerate(A):
+        t_pos[i, : len(a["t_pos"])] = torch.tensor(a["t_pos"], dtype=torch.long)
+    dec = torch.tensor([[J if (l == len(a["t_pos"])) else l for l in a["dec_labels"]] for a in A])
+    len_idx, len_labels = [], []
+    for i, a in enumerate(A):
+        for (ci, j, k, b) in a["len_items"]:
+            len_idx.append((i, a["call_pos"][ci], a["arg_pos"][(j, k)]))
+            len_labels.append(b)
+    batch = {"a_ids": a_ids, "a_attn": a_attn, "t_pos": t_pos, "t_mask": t_pos >= 0,
+             "call_pos": torch.tensor([a["call_pos"] for a in A]), "dec_labels": dec,
+             "len_idx": torch.tensor(len_idx, dtype=torch.long).reshape(-1, 3),
+             "len_labels": torch.tensor(len_labels, dtype=torch.long)}
+
+    # pass B with a random mask ratio per example (masked-diffusion objective)
+    B = [it["b"] for it in items if it["b"] is not None and (it["b"]["slot_pos"] or it["b"].get("opt_slots"))]
+    if B:
+        seqs, sb, sp, sl = [], [], [], []
+        for bi, b in enumerate(B):
+            ids, targets = _mask_slots(b, rng, v3_mask)
+            for p, t in targets:
+                sb.append(bi); sp.append(p); sl.append(t)
+            seqs.append(ids)
+        b_ids, b_attn = _pad(seqs, tok.pad_token_id)
+        batch.update({"b_ids": b_ids, "b_attn": b_attn, "slot_b": torch.tensor(sb), "slot_p": torch.tensor(sp),
+                      "slot_labels": torch.tensor(sl)})
+    else:
+        batch["b_ids"] = None
+    if "q_pos" in A[0]:
+        from .v3feat import v3_extras
+        batch.update(v3_extras(A, list(enumerate(B))))
+    return batch
+
+
+def _collate_cached(items, tok, rng, v3_mask="clean"):
+    """One sequence per example: pass-A context + slot block, read with the block-causal mask."""
+    A, seqs, ctx, sb, sp, sl = [], [], [], [], [], []
+    for i, it in enumerate(items):
+        a, b = it["a"], it["b"]
+        A.append(a)
+        ctx.append(len(a["input_ids"]))
+        if b is not None and (b["slot_pos"] or b.get("opt_slots")):
+            ids, targets = _mask_slots(b, rng, v3_mask)
+            for p, t in targets:
+                sb.append(i); sp.append(p); sl.append(t)
+            seqs.append(ids)
+        else:
+            seqs.append(a["input_ids"])
+    ids, attn = _pad(seqs, tok.pad_token_id)
+    J = max(max(len(a["t_pos"]) for a in A), 1)
+    t_pos = torch.full((len(A), J), -1, dtype=torch.long)
+    for i, a in enumerate(A):
+        t_pos[i, : len(a["t_pos"])] = torch.tensor(a["t_pos"], dtype=torch.long)
+    dec = torch.tensor([[J if (l == len(a["t_pos"])) else l for l in a["dec_labels"]] for a in A])
+    len_idx, len_labels = [], []
+    for i, a in enumerate(A):
+        for (ci, j, k, b) in a["len_items"]:
+            len_idx.append((i, a["call_pos"][ci], a["arg_pos"][(j, k)]))
+            len_labels.append(b)
+    out = {"a_ids": ids, "a_attn": attn, "t_pos": t_pos, "t_mask": t_pos >= 0,
+            "call_pos": torch.tensor([a["call_pos"] for a in A]), "dec_labels": dec,
+            "len_idx": torch.tensor(len_idx, dtype=torch.long).reshape(-1, 3),
+            "len_labels": torch.tensor(len_labels, dtype=torch.long),
+            "ctx_len": torch.tensor(ctx), "b_ids": None, "slot_b": torch.tensor(sb, dtype=torch.long),
+            "slot_p": torch.tensor(sp, dtype=torch.long), "slot_labels": torch.tensor(sl, dtype=torch.long)}
+    if "q_pos" in A[0]:
+        from .v3feat import v3_extras
+        out.update(v3_extras(A, [(i, it["b"]) for i, it in enumerate(items)
+                                 if it["b"] is not None and (it["b"]["slot_pos"] or it["b"].get("opt_slots"))]))
+    return out
+
+
+def token_batches(lengths: list[int], tokens_per_batch: int, rng: random.Random, max_bs: int = 256):
+    """Length-bucketed batches (sorted by length, chunked by token budget, shuffled)."""
+    order = sorted(range(len(lengths)), key=lambda i: lengths[i])
+    batches, cur, cur_max = [], [], 0
+    for i in order:
+        m = max(cur_max, lengths[i])
+        if cur and (m * (len(cur) + 1) > tokens_per_batch or len(cur) >= max_bs):
+            batches.append(cur); cur, m = [], lengths[i]
+        cur.append(i); cur_max = m
+    if cur:
+        batches.append(cur)
+    rng.shuffle(batches)
+    return batches
+
+
+def hard_negatives(examples: list[Example], n: int, rng: random.Random) -> list[Example]:
+    """Same query, same domain: drop every tool the gold answer uses, keep its related neighbours -> no call."""
+    out = []
+    pool = [e for e in examples if e.calls and len(e.tools) > len({c.name for c in e.calls})]
+    rng.shuffle(pool)
+    for e in pool[:n]:
+        needed = {c.name for c in e.calls}
+        tools = [t for t in e.tools if t.name not in needed]
+        out.append(Example(e.query, tools, [], id=f"hardneg_{len(out)}", source="xlam_hardneg", gold=[]))
+    return out
+
+
+def irrelevance_negatives(examples: list[Example], n: int, rng: random.Random) -> list[Example]:
+    """Pair real queries with tool lists from other examples that don't contain any needed tool."""
+    out = []
+    pool = [e for e in examples if e.tools]
+    while len(out) < n:
+        q, other = rng.choice(pool), rng.choice(pool)
+        needed = {c.name for c in q.calls}
+        tools = [t for t in other.tools if t.name not in needed]
+        if not tools or other is q:
+            continue
+        out.append(Example(q.query, tools, [], id=f"neg_{len(out)}", source="xlam_neg", gold=[]))
+    return out

code/actionman/features.py

@@ -0,0 +1,129 @@
+"""Tokenizer setup and the two sequence builders.
+
+Pass A:  [CLS] [T] sig_1 [T] sig_2 ... [USER] query [CALL] x4 [SEP]
+Pass B:  <pass A without [SEP]> [ARGS] tool( p= [MASK]*n , q= [MASK]*m ) ... [SEP]
+"""
+from __future__ import annotations
+
+from .schema import Example, signature_pieces, value_to_text
+
+SPECIAL = ["[T]", "[CALL]", "[ARGS]", "[USER]", "[OPT]"]   # [OPT]: v3 option marker (appended, ids of the rest unchanged)
+MAX_CALLS = 4
+BUCKETS = [0, 4, 8, 16, 32]          # slot counts for bucket ids 0..4; id 5 = overflow
+OVERFLOW = 5
+N_BUCKETS = 6
+MAX_LEN_A = 896
+MAX_LEN_B = 1024
+
+
+def load_tokenizer(name: str = "answerdotai/ModernBERT-large"):
+    from transformers import AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(name)
+    tok.add_special_tokens({"additional_special_tokens": SPECIAL})
+    tok.am = {s: tok.convert_tokens_to_ids(s) for s in SPECIAL}
+    return tok
+
+
+def bucket_of(n_tokens: int) -> int:
+    for b in range(1, 5):
+        if n_tokens <= BUCKETS[b]:
+            return b
+    return OVERFLOW
+
+
+def _enc(tok, text: str) -> list[int]:
+    return tok(text, add_special_tokens=False).input_ids
+
+
+def value_tokens(tok, value) -> list[int]:
+    return _enc(tok, value_to_text(value))
+
+
+def build_pass_a(tok, ex: Example, desc_words: int = 8, max_len: int = MAX_LEN_A) -> dict | None:
+    """Context sequence + decision / length labels. Returns None if it can't fit."""
+    for dw in (desc_words, 0):
+        ids = [tok.cls_token_id]
+        t_pos, arg_pos = [], {}
+        for j, tool in enumerate(ex.tools):
+            t_pos.append(len(ids))
+            ids.append(tok.am["[T]"])
+            for text, kind, k in signature_pieces(tool, dw):
+                if kind == "param":
+                    arg_pos[(j, k)] = len(ids)
+                ids += _enc(tok, text)
+        ids.append(tok.am["[USER]"])
+        ids += _enc(tok, ex.query)
+        call_pos = list(range(len(ids), len(ids) + MAX_CALLS))
+        ids += [tok.am["[CALL]"]] * MAX_CALLS + [tok.sep_token_id]
+        if len(ids) <= max_len:
+            break
+    else:
+        return None
+
+    names = [t.name for t in ex.tools]
+    dec_labels, len_items = [], []
+    for i in range(MAX_CALLS):
+        if i < len(ex.calls):
+            c = ex.calls[i]
+            if c.name not in names:
+                return None
+            j = names.index(c.name)
+            dec_labels.append(j)
+            for k, p in enumerate(ex.tools[j].params):
+                b = bucket_of(len(value_tokens(tok, c.args[p.name]))) if p.name in c.args else 0
+                len_items.append((i, j, k, b))
+        else:
+            dec_labels.append(len(ex.tools))  # NONE
+    if len(ex.calls) > MAX_CALLS:
+        dec_labels = [-100] * MAX_CALLS     # can't represent; train nothing on decisions
+        len_items = [it for it in len_items if it[0] < MAX_CALLS]
+    return {"input_ids": ids, "t_pos": t_pos, "call_pos": call_pos, "arg_pos": arg_pos,
+            "dec_labels": dec_labels, "len_items": len_items, "desc_words": dw}
+
+
+def build_pass_b(tok, ex: Example, a: dict, plan: list[tuple[int, dict[int, int]]],
+                 with_targets: bool, max_len: int = MAX_LEN_B, keep_context: bool = False,
+                 markers: bool = False) -> dict | None:
+    """Append the slot block for `plan` = [(tool_j, {param_k: bucket})] per call.
+    Buckets 1..4 get BUCKETS[b] mask slots; 0 (absent) and 5 (overflow) get none."""
+    # keep_context=True (cached variant): the full pass-A sequence is the prefix, so its states can be reused
+    ids = (a["input_ids"] if keep_context else a["input_ids"][:-1]) + [tok.am["[ARGS]"]]
+    ctx_len = len(a["input_ids"]) if keep_context else len(ids) - 1
+    slots, targets, index = [], [], []
+    for i, (j, buckets) in enumerate(plan):
+        tool = ex.tools[j]
+        # markers=True (v2): '#1 tool(' '#2 tool(' so repeated calls of the same tool are distinguishable
+        ids += _enc(tok, (f" #{i + 1} " if markers else "") + tool.name + "(")
+        first = True
+        for k, p in enumerate(tool.params):
+            b = buckets.get(k, 0)
+            if b == 0 or b == OVERFLOW:
+                continue
+            ids += _enc(tok, ("" if first else ", ") + p.name + "=")
+            first = False
+            n = BUCKETS[b]
+            pos = list(range(len(ids), len(ids) + n))
+            ids += [tok.mask_token_id] * n
+            slots.append(pos)
+            index.append((i, k))
+            if with_targets:
+                v = value_tokens(tok, ex.calls[i].args[p.name])[:n]
+                targets.append(v + [tok.pad_token_id] * (n - len(v)))
+        ids += _enc(tok, ")")
+    ids.append(tok.sep_token_id)
+    if len(ids) > max_len:
+        return None
+    out = {"input_ids": ids, "slot_pos": slots, "slot_index": index, "ctx_len": ctx_len}
+    if with_targets:
+        out["slot_targets"] = targets
+    return out
+
+
+def gold_plan(a: dict, ex: Example) -> list[tuple[int, dict[int, int]]]:
+    """One (tool_j, {param_k: bucket}) entry per gold call, in call order - including calls to tools
+    with no parameters (an empty dict), so plan index == call index."""
+    plan = []
+    for i in range(min(len(ex.calls), MAX_CALLS)):
+        buckets = {k: b for (ci, _, k, b) in a["len_items"] if ci == i}
+        plan.append((a["dec_labels"][i], buckets))
+    return plan

code/actionman/infer.py

@@ -0,0 +1,132 @@
+"""Inference: pass A (decide tools + arg lengths) -> slot block -> mask-predict refinement."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+from .features import BUCKETS, MAX_CALLS, OVERFLOW, SPECIAL, build_pass_a, build_pass_b
+from .schema import Call, Example, text_to_value
+
+
+@dataclass
+class PredCall:
+    name: str
+    args: dict = field(default_factory=dict)
+    conf_tool: float = 0.0
+    conf_args: dict = field(default_factory=dict)
+    escalated: list = field(default_factory=list)
+
+    def as_call(self) -> Call:
+        return Call(self.name, self.args)
+
+
+@dataclass
+class Prediction:
+    calls: list[PredCall]
+    passes: int = 0
+    skipped: bool = False
+
+
+def _banned_ids(tok, vocab: int) -> torch.Tensor:
+    ids = [tok.mask_token_id, tok.cls_token_id, tok.sep_token_id] + [tok.am[s] for s in SPECIAL]
+    ids += list(range(len(tok), vocab))
+    return torch.tensor(sorted(set(i for i in ids if i < vocab)))
+
+
+@torch.no_grad()
+def predict(model, tok, ex: Example, T: int = 2, tau: float = 0.85, device="cpu", cached: bool = False,
+            markers: bool = False, v3: bool = False, snap: bool = False, require: bool = False) -> Prediction:
+    if v3:
+        from .infer_v3 import predict_v3
+        return predict_v3(model, tok, ex, T=T, tau=tau, device=device, cached=cached, markers=markers, snap=snap, require=require)
+    model.eval()
+    ex0 = Example(ex.query, ex.tools, [], ex.id, ex.source)
+    a = build_pass_a(tok, ex0)
+    if a is None or not ex.tools:
+        return Prediction([], 0, skipped=a is None)
+    ids = torch.tensor([a["input_ids"]], device=device)
+    if cached:
+        from .cached import encode_context, encode_slots
+        h, kv = encode_context(model.mlm.model, ids)
+    else:
+        h = model.encode(ids, torch.ones_like(ids))
+    J = len(a["t_pos"])
+    call_pos = torch.tensor([a["call_pos"]], device=device)
+    t_pos = torch.tensor([a["t_pos"]], device=device)
+    probs = model.decision_logits(h, call_pos, t_pos, torch.ones_like(t_pos, dtype=torch.bool)).float().softmax(-1)[0]
+
+    chosen = []                                   # (call_i, tool_j, conf)
+    for i in range(MAX_CALLS):
+        j = int(probs[i].argmax())
+        if j == J:
+            break
+        chosen.append((i, j, float(probs[i, j])))
+    if not chosen:
+        return Prediction([], 1)
+
+    rows, keys = [], []
+    for i, j, _ in chosen:
+        for k in range(len(ex.tools[j].params)):
+            rows.append((0, a["call_pos"][i], a["arg_pos"][(j, k)]))
+            keys.append((i, j, k))
+    len_logits = model.length_logits(h, torch.tensor(rows, device=device).reshape(-1, 3)).float()
+    plan, preds = [], []
+    for ci, (i, j, conf) in enumerate(chosen):
+        buckets, pc = {}, PredCall(ex.tools[j].name, conf_tool=conf)
+        for r, (ri, rj, k) in enumerate(keys):
+            if ri != i:
+                continue
+            lg = len_logits[r].clone()
+            if not ex.tools[j].params[k].optional:
+                lg[0] = float("-inf")            # required params can't be absent
+            b = int(lg.argmax())
+            buckets[k] = b
+            if b == OVERFLOW:
+                pc.escalated.append(ex.tools[j].params[k].name)
+        plan.append((j, buckets))
+        preds.append(pc)
+
+    b = build_pass_b(tok, ex0, a, plan, with_targets=False, keep_context=cached, markers=markers)
+    passes = 1
+    if b is not None and b["slot_pos"]:
+        ids_b = torch.tensor([b["input_ids"]], device=device)
+        flat = [p for pos in b["slot_pos"] for p in pos]
+        fpos = torch.tensor(flat, device=device)
+        conf = torch.zeros(len(flat), device=device)
+        masked = torch.ones(len(flat), dtype=torch.bool, device=device)
+        vocab = model.mlm.decoder.weight.shape[0]
+        banned = _banned_ids(tok, vocab).to(device)
+        for t in range(T):
+            passes = t + 1
+            if cached:
+                c0 = b["ctx_len"]
+                hs = encode_slots(model.mlm.model, ids_b[:, c0:], c0, kv)[0, fpos - c0]
+                logits = model.mlm.decoder(model.mlm.head(hs)).float()
+            else:
+                logits = model.forward_b(ids_b, torch.ones_like(ids_b), torch.zeros_like(fpos), fpos).float()
+            logits[:, banned] = float("-inf")
+            p, tokens = logits.softmax(-1).max(-1)
+            conf = torch.where(masked, p, conf)
+            ids_b[0, fpos[masked]] = tokens[masked]
+            if t == T - 1:
+                break
+            masked = conf < tau
+            if not masked.any():
+                break
+            ids_b[0, fpos[masked]] = tok.mask_token_id
+        out_tokens = ids_b[0, fpos].tolist()
+        confs = conf.tolist()
+        cursor = 0
+        for (i, k), pos in zip(b["slot_index"], b["slot_pos"]):
+            n = len(pos)
+            toks, cs = out_tokens[cursor: cursor + n], confs[cursor: cursor + n]
+            cursor += n
+            if tok.pad_token_id in toks:
+                toks = toks[: toks.index(tok.pad_token_id)]
+            text = tok.decode(toks, skip_special_tokens=True).strip()
+            ci = [c[0] for c in chosen].index(i)
+            param = ex.tools[chosen[ci][1]].params[k]
+            preds[ci].args[param.name] = text_to_value(text, param.type)
+            preds[ci].conf_args[param.name] = float(min(cs)) if cs else 1.0
+    return Prediction(preds, passes + 1)

code/actionman/infer_v3.py

@@ -0,0 +1,182 @@
+"""v3 inference: pass A (tools, lengths, list counts) -> slot block with option / item slots ->
+copy-or-generate per value (decided on the first pass) -> mask-predict refinement of the generated slots."""
+from __future__ import annotations
+
+from collections import defaultdict
+
+import torch
+
+from .features import MAX_CALLS, OVERFLOW
+from .infer import Prediction, PredCall, _banned_ids
+from .schema import Example, param_options, text_to_value
+from .span import best_span, snap_word
+from .v3feat import build_pass_a_v3, build_pass_b_v3
+
+
+@torch.no_grad()
+def predict_v3(model, tok, ex: Example, T: int = 2, tau: float = 0.85, device="cpu", cached: bool = False,
+               markers: bool = True, snap: bool = False, require: bool = False) -> Prediction:
+    """snap: widen copied spans to whole words/numbers (inference-only fix for copies cut mid-word).
+    require: like tool_choice="required" -- the first call never abstains; it takes the best tool instead."""
+    model.eval()
+    ex0 = Example(ex.query, ex.tools, [], ex.id, ex.source)
+    a = build_pass_a_v3(tok, ex0)
+    if a is None or not ex.tools:
+        return Prediction([], 0, skipped=a is None)
+    ids = torch.tensor([a["input_ids"]], device=device)
+    if cached:
+        from .cached import encode_context, encode_slots
+        h, kv = encode_context(model.mlm.model, ids)
+    else:
+        h = model.encode(ids, torch.ones_like(ids))
+    J = len(a["t_pos"])
+    call_pos = torch.tensor([a["call_pos"]], device=device)
+    t_pos = torch.tensor([a["t_pos"]], device=device)
+    probs = model.decision_logits(h, call_pos, t_pos, torch.ones_like(t_pos, dtype=torch.bool)).float().softmax(-1)[0]
+    chosen = []
+    for i in range(MAX_CALLS):
+        j = int(probs[i].argmax())
+        if j == J and require and i == 0:
+            j = int(probs[i, :J].argmax())
+        if j == J:
+            break
+        chosen.append((i, j, float(probs[i, j])))
+    if not chosen:
+        return Prediction([], 1)
+
+    rows, keys = [], []
+    for i, j, _ in chosen:
+        for k in range(len(ex.tools[j].params)):
+            rows.append((0, a["call_pos"][i], a["arg_pos"][(j, k)]))
+            keys.append((i, j, k))
+    idx = torch.tensor(rows, device=device).reshape(-1, 3)
+    len_logits = model.length_logits(h, idx).float()
+    cnt_logits = model.count_logits(h, idx).float()
+    plan, preds = [], []
+    for i, j, conf in chosen:
+        modes, pc = {}, PredCall(ex.tools[j].name, conf_tool=conf)
+        for r, (ri, _, k) in enumerate(keys):
+            if ri != i:
+                continue
+            p = ex.tools[j].params[k]
+            lg = len_logits[r].clone()
+            if not p.optional:
+                lg[0] = float("-inf")
+            b = int(lg.argmax())
+            if b == 0:
+                continue
+            if param_options(p) is not None:
+                modes[k] = ("opt", 0)
+            elif p.type == "list" and int(cnt_logits[r].argmax()) >= 1:
+                modes[k] = ("list", int(cnt_logits[r].argmax()))
+            else:
+                modes[k] = ("text", b)
+                if b == OVERFLOW:
+                    pc.escalated.append(p.name)
+        plan.append((j, modes))
+        preds.append(pc)
+
+    bb = build_pass_b_v3(tok, ex0, a, plan, with_targets=False, keep_context=cached, markers=markers)
+    if bb is None or not (bb["slot_pos"] or bb["opt_slots"]):
+        return Prediction(preds, 2)
+    ids_b = torch.tensor([bb["input_ids"]], device=device)
+    c0 = bb["ctx_len"]
+
+    def forward():
+        if cached:
+            hs = encode_slots(model.mlm.model, ids_b[:, c0:], c0, kv)
+            return torch.cat([h, hs], dim=1)
+        return model.encode(ids_b, torch.ones_like(ids_b))
+
+    flat = [p for pos in bb["slot_pos"] for p in pos]
+    fpos = torch.tensor(flat, dtype=torch.long, device=device)
+    conf = torch.zeros(len(flat), device=device)
+    masked = torch.ones(len(flat), dtype=torch.bool, device=device)
+    fixed = torch.zeros(len(flat), dtype=torch.bool, device=device)
+    vocab = model.mlm.decoder.weight.shape[0]
+    banned = _banned_ids(tok, vocab).to(device)
+    copied = {}                                     # value index -> (text, conf)
+    starts = [0]
+    for pos in bb["slot_pos"]:
+        starts.append(starts[-1] + len(pos))
+    passes, hb = 0, None
+    for t in range(T):
+        passes = t + 1
+        hb = forward()
+        if t == 0 and bb["anchors"]:
+            ab = torch.zeros(len(bb["anchors"]), dtype=torch.long, device=device)
+            ap = torch.tensor(bb["anchors"], device=device)
+            qpos = torch.tensor([a["q_pos"]], device=device).expand(len(bb["anchors"]), -1)
+            g, st, en = model.copy_logits(hb, ab, ap, qpos)
+            g = g.float().sigmoid()
+            for v in range(len(bb["anchors"])):
+                if g[v] <= 0.5 or not a["q_pos"]:
+                    continue
+                s, e, pspan = best_span(st[v], en[v])
+                cs, ce = a["q_offs"][s][0], a["q_offs"][e][1]
+                if snap:
+                    cs, ce = snap_word(ex.query, cs, ce)
+                text = ex.query[cs: ce].strip()
+                if not text:
+                    continue
+                copied[v] = (text, float(g[v]) * pspan)
+                n = len(bb["slot_pos"][v])
+                span_ids = [a["input_ids"][q] for q in a["q_pos"][s: e + 1]][:n]
+                span_ids += [tok.pad_token_id] * (n - len(span_ids))
+                sl = slice(starts[v], starts[v + 1])
+                ids_b[0, fpos[sl]] = torch.tensor(span_ids, device=device)
+                fixed[sl] = True
+        if not flat:
+            break
+        logits = model.mlm.decoder(model.mlm.head(hb[0, fpos])).float()
+        logits[:, banned] = float("-inf")
+        p, tokens = logits.softmax(-1).max(-1)
+        upd = masked & ~fixed
+        conf = torch.where(upd, p, conf)
+        ids_b[0, fpos[upd]] = tokens[upd]
+        conf = torch.where(fixed, torch.ones_like(conf), conf)
+        if t == T - 1:
+            break
+        masked = (conf < tau) & ~fixed
+        if not masked.any():
+            break
+        ids_b[0, fpos[masked]] = tok.mask_token_id
+
+    ci_of = {c[0]: n for n, c in enumerate(chosen)}
+    out_tokens, confs = ids_b[0, fpos].tolist() if flat else [], conf.tolist()
+    items = defaultdict(list)                       # (call, param) -> [(text, conf)] for list items
+    for v, ((i, k, item), pos) in enumerate(zip(bb["slot_index"], bb["slot_pos"])):
+        if v in copied:
+            text, cv = copied[v]
+        else:
+            toks = out_tokens[starts[v]: starts[v + 1]]
+            if tok.pad_token_id in toks:
+                toks = toks[: toks.index(tok.pad_token_id)]
+            text = tok.decode(toks, skip_special_tokens=True).strip()
+            cs = confs[starts[v]: starts[v + 1]]
+            cv = float(min(cs)) if cs else 1.0
+        ci = ci_of[i]
+        param = ex.tools[chosen[ci][1]].params[k]
+        if item is None:
+            preds[ci].args[param.name] = text_to_value(text, param.type)
+            preds[ci].conf_args[param.name] = cv
+        else:
+            items[(ci, param.name)].append((text, cv))
+    for (ci, name), vals in items.items():
+        preds[ci].args[name] = [text_to_value(t, "any") for t, _ in vals]
+        preds[ci].conf_args[name] = min(c for _, c in vals)
+    if bb["opt_slots"]:
+        ob = torch.zeros(len(bb["opt_slots"]), dtype=torch.long, device=device)
+        oa = torch.tensor([s[2] for s in bb["opt_slots"]], device=device)
+        M = max(len(s[3]) for s in bb["opt_slots"])
+        opos = torch.full((len(bb["opt_slots"]), M), -1, dtype=torch.long, device=device)
+        for n, s in enumerate(bb["opt_slots"]):
+            opos[n, : len(s[3])] = torch.tensor(s[3], device=device)
+        po = model.option_logits(hb, ob, oa, opos).float().softmax(-1)
+        for n, (i, k, _, _, _) in enumerate(bb["opt_slots"]):
+            ci = ci_of[i]
+            param = ex.tools[chosen[ci][1]].params[k]
+            m = int(po[n].argmax())
+            preds[ci].args[param.name] = param_options(param)[m]
+            preds[ci].conf_args[param.name] = float(po[n, m])
+    return Prediction(preds, passes + 1)

code/actionman/model.py

@@ -0,0 +1,176 @@
+"""ActionMan: ModernBERT encoder + MLM head (slot filler) + decision head + length head."""
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .features import N_BUCKETS
+
+
+class ActionMan(nn.Module):
+    def __init__(self, mlm, v3: bool = False):
+        super().__init__()
+        self.mlm = mlm
+        d = mlm.config.hidden_size
+        self.q = nn.Linear(d, d)
+        self.k = nn.Linear(d, d)
+        self.none = nn.Linear(d, 1)
+        self.len_head = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, N_BUCKETS))
+        self.scale = 1.0 / math.sqrt(d)
+        self.v3 = v3
+        if v3:   # copy switch + span pointer, option picker, list item count
+            from .v3feat import N_COUNT
+            self.gate = nn.Linear(d, 1)
+            self.sq, self.sk, self.eq, self.ek = nn.Linear(d, d), nn.Linear(d, d), nn.Linear(d, d), nn.Linear(d, d)
+            self.opt_q, self.opt_k = nn.Linear(d, d), nn.Linear(d, d)
+            self.count_head = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, N_COUNT))
+
+    @classmethod
+    def from_encoder(cls, name_or_config, tok, attn_impl: str | None = None, v3: bool = False):
+        from transformers import ModernBertConfig, ModernBertForMaskedLM
+        if isinstance(name_or_config, ModernBertConfig):
+            mlm = ModernBertForMaskedLM(name_or_config)
+        else:
+            kw = {"attn_implementation": attn_impl} if attn_impl else {}
+            mlm = ModernBertForMaskedLM.from_pretrained(name_or_config, **kw)
+        mlm.resize_token_embeddings(len(tok), pad_to_multiple_of=64, mean_resizing=False)
+        dec = mlm.decoder
+        if dec.bias is not None and dec.bias.shape[0] != dec.weight.shape[0]:
+            bias = torch.zeros(dec.weight.shape[0], dtype=dec.bias.dtype)
+            bias[: dec.bias.shape[0]] = dec.bias.data
+            dec.bias = nn.Parameter(bias)
+            dec.out_features = dec.weight.shape[0]
+        return cls(mlm, v3=v3)
+
+    def head_params(self):
+        return [p for n, p in self.named_parameters() if not n.startswith("mlm.")]
+
+    def encode(self, ids, attn):
+        return self.mlm.model(input_ids=ids, attention_mask=attn).last_hidden_state
+
+    def decision_logits(self, h, call_pos, t_pos, t_mask):
+        bi = torch.arange(h.shape[0], device=h.device)[:, None]
+        call_h = h[bi, call_pos]                                               # [B,4,d]
+        t_h = h[bi, t_pos.clamp(min=0)]                                        # [B,J,d]
+        logits = torch.einsum("bid,bjd->bij", self.q(call_h), self.k(t_h)) * self.scale
+        logits = logits.masked_fill(~t_mask[:, None, :], float("-inf"))
+        return torch.cat([logits, self.none(call_h)], dim=-1)                 # [B,4,J+1]; NONE = J
+
+    def length_logits(self, h, len_idx):
+        """len_idx: [N,3] rows of (batch index, [CALL] position, param-name position)."""
+        if not len_idx.numel():
+            return h.new_zeros(0, N_BUCKETS)
+        hc, ha = h[len_idx[:, 0], len_idx[:, 1]], h[len_idx[:, 0], len_idx[:, 2]]
+        return self.len_head(torch.cat([hc, ha, hc * ha], dim=-1))
+
+    def forward_a(self, batch):
+        h = self.encode(batch["a_ids"], batch["a_attn"])                     # [B,L,d]
+        dec = self.decision_logits(h, batch["call_pos"], batch["t_pos"], batch["t_mask"])
+        return dec, self.length_logits(h, batch["len_idx"])
+
+    def forward_b(self, ids, attn, slot_b, slot_p):
+        h = self.encode(ids, attn)
+        hs = h[slot_b, slot_p]
+        return self.mlm.decoder(self.mlm.head(hs))                            # [M,V]
+
+    # ---------- v3 heads ----------
+
+    def copy_logits(self, h, ab, ap, qpos):
+        """Anchors (batch ab, pos ap) [N]; qpos [N,Q] query positions (-1 = pad). -> gate [N], start/end [N,Q]."""
+        ha = h[ab, ap]
+        hq = h[ab[:, None], qpos.clamp(min=0)]
+        pad = qpos < 0
+        st = torch.einsum("nd,nqd->nq", self.sq(ha), self.sk(hq)) * self.scale
+        en = torch.einsum("nd,nqd->nq", self.eq(ha), self.ek(hq)) * self.scale
+        return self.gate(ha).squeeze(-1), st.masked_fill(pad, -1e4), en.masked_fill(pad, -1e4)
+
+    def option_logits(self, h, ob, oa, opos):
+        """Option anchors (ob, oa) [N]; opos [N,M] positions of their [OPT] markers (-1 = pad) -> [N,M]."""
+        ho = h[ob[:, None], opos.clamp(min=0)]
+        lg = torch.einsum("nd,nmd->nm", self.opt_q(h[ob, oa]), self.opt_k(ho)) * self.scale
+        return lg.masked_fill(opos < 0, -1e4)
+
+    def count_logits(self, h, idx):
+        if not idx.numel():
+            return h.new_zeros(0, self.count_head[-1].out_features)
+        hc, ha = h[idx[:, 0], idx[:, 1]], h[idx[:, 0], idx[:, 2]]
+        return self.count_head(torch.cat([hc, ha, hc * ha], dim=-1))
+
+    def v3_losses(self, h_b, h_a, batch):
+        z = h_a.sum() * 0
+        l_gate = l_span = l_opt = l_cnt = z
+        if batch["anchor_b"].numel():
+            g, st, en = self.copy_logits(h_b, batch["anchor_b"], batch["anchor_p"], batch["qpos"][batch["anchor_b"]])
+            l_gate = F.binary_cross_entropy_with_logits(g.float(), batch["gate_y"])
+            sel = batch["span_s"] >= 0
+            if sel.any():
+                l_span = 0.5 * (F.cross_entropy(st[sel].float(), batch["span_s"][sel]) +
+                                F.cross_entropy(en[sel].float(), batch["span_e"][sel]))
+        valid = batch["opt_y"] >= 0
+        if valid.any():
+            lo = self.option_logits(h_b, batch["opt_b"][valid], batch["opt_anchor"][valid], batch["opt_pos"][valid])
+            l_opt = F.cross_entropy(lo.float(), batch["opt_y"][valid])
+        if batch["count_idx"].numel():
+            l_cnt = F.cross_entropy(self.count_logits(h_a, batch["count_idx"]).float(), batch["count_y"])
+        parts = {"gate": l_gate.item(), "span": l_span.item(), "opt": l_opt.item(), "count": l_cnt.item()}
+        return l_gate + l_span + l_opt + l_cnt, parts
+
+    def loss_v3(self, batch):
+        """Standard-mask v3: pass A and pass B are separate forwards."""
+        h_a = self.encode(batch["a_ids"], batch["a_attn"])
+        dec_logits = self.decision_logits(h_a, batch["call_pos"], batch["t_pos"], batch["t_mask"])
+        len_logits = self.length_logits(h_a, batch["len_idx"])
+        l_dec = F.cross_entropy(dec_logits.flatten(0, 1).float(), batch["dec_labels"].flatten(), ignore_index=-100)
+        l_len = F.cross_entropy(len_logits.float(), batch["len_labels"]) if len_logits.shape[0] else h_a.sum() * 0
+        l_slot = h_a.sum() * 0
+        h_b = h_a
+        if batch.get("b_ids") is not None:
+            h_b = self.encode(batch["b_ids"], batch["b_attn"])
+            if batch["slot_b"].numel():
+                hs = h_b[batch["slot_b"], batch["slot_p"]]
+                l_slot = F.cross_entropy(self.mlm.decoder(self.mlm.head(hs)).float(), batch["slot_labels"])
+        if batch.get("b_ids") is None:
+            batch = {**batch, "anchor_b": batch["anchor_b"][:0], "opt_y": batch["opt_y"][:0]}
+        l_v3, p3 = self.v3_losses(h_b, h_a, batch)
+        return l_dec + l_len + l_slot + l_v3, {"dec": l_dec.item(), "len": l_len.item(), "slot": l_slot.item(), **p3}
+
+    def loss(self, batch):
+        if batch.get("ctx_len") is not None:
+            return self.loss_cached(batch)
+        if self.v3:
+            return self.loss_v3(batch)
+        dec_logits, len_logits = self.forward_a(batch)
+        l_dec = F.cross_entropy(dec_logits.flatten(0, 1).float(), batch["dec_labels"].flatten(), ignore_index=-100)
+        l_len = (F.cross_entropy(len_logits.float(), batch["len_labels"]) if len_logits.shape[0]
+                 else dec_logits.sum() * 0)
+        if batch.get("b_ids") is not None and batch["slot_b"].numel():
+            slot_logits = self.forward_b(batch["b_ids"], batch["b_attn"], batch["slot_b"], batch["slot_p"])
+            l_slot = F.cross_entropy(slot_logits.float(), batch["slot_labels"])
+        else:
+            l_slot = dec_logits.sum() * 0
+        total = l_dec + l_len + l_slot
+        return total, {"dec": l_dec.item(), "len": l_len.item(), "slot": l_slot.item()}
+
+    def loss_cached(self, batch):
+        """Block-causal variant: one forward; heads read context states, which never see the slots."""
+        from .cached import block_masks
+        masks = block_masks(batch["a_attn"], batch["ctx_len"], self.mlm.config.sliding_window)
+        h = self.mlm.model(input_ids=batch["a_ids"], attention_mask=masks).last_hidden_state
+        dec_logits = self.decision_logits(h, batch["call_pos"], batch["t_pos"], batch["t_mask"])
+        len_logits = self.length_logits(h, batch["len_idx"])
+        l_dec = F.cross_entropy(dec_logits.flatten(0, 1).float(), batch["dec_labels"].flatten(), ignore_index=-100)
+        l_len = F.cross_entropy(len_logits.float(), batch["len_labels"]) if len_logits.shape[0] else h.sum() * 0
+        if batch["slot_b"].numel():
+            hs = h[batch["slot_b"], batch["slot_p"]]
+            l_slot = F.cross_entropy(self.mlm.decoder(self.mlm.head(hs)).float(), batch["slot_labels"])
+        else:
+            l_slot = h.sum() * 0
+        parts = {"dec": l_dec.item(), "len": l_len.item(), "slot": l_slot.item()}
+        total = l_dec + l_len + l_slot
+        if self.v3:
+            l_v3, p3 = self.v3_losses(h, h, batch)
+            total, parts = total + l_v3, {**parts, **p3}
+        return total, parts

code/actionman/schema.py

@@ -0,0 +1,252 @@
+"""Unified tool/call schema, converters from xLAM and BFCL, compact signatures, value <-> text."""
+from __future__ import annotations
+
+import ast
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+TYPES = ("str", "int", "float", "bool", "list", "dict", "any")
+
+
+@dataclass
+class Param:
+    name: str
+    type: str = "any"
+    optional: bool = False
+    default: Any = None
+    description: str = ""
+    options: list | None = None          # declared enum values (v3 option slots)
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str = ""
+    params: list[Param] = field(default_factory=list)
+
+
+@dataclass
+class Call:
+    name: str
+    args: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Example:
+    query: str
+    tools: list[Tool]
+    calls: list[Call]
+    id: str = ""
+    source: str = ""
+    # gold in BFCL "possible answer" form: [{func: {arg: [acceptable values]}}]; "" = may be omitted
+    gold: list[dict] = field(default_factory=list)
+
+
+# ---------- type normalisation ----------
+
+def norm_type(t: Any) -> tuple[str, bool]:
+    """Map xLAM ('str, optional', 'List[int]') or JSON-schema ('integer') types to (TYPE, optional)."""
+    s = str(t or "any").strip()
+    low = s.lower()
+    optional = "optional" in low
+    head = low.split(",")[0].split("[")[0].strip()
+    table = {
+        "str": "str", "string": "str",
+        "int": "int", "integer": "int",
+        "float": "float", "number": "float", "double": "float",
+        "bool": "bool", "boolean": "bool",
+        "list": "list", "array": "list", "tuple": "list", "set": "list",
+        "dict": "dict", "object": "dict",
+    }
+    return table.get(head, "any"), optional
+
+
+# ---------- converters ----------
+
+def from_xlam(row: dict) -> Example:
+    tools_raw = json.loads(row["tools"]) if isinstance(row["tools"], str) else row["tools"]
+    answers = json.loads(row["answers"]) if isinstance(row["answers"], str) else row["answers"]
+    tools = []
+    for t in tools_raw:
+        params = []
+        for pname, spec in (t.get("parameters") or {}).items():
+            spec = spec or {}
+            ty, opt = norm_type(spec.get("type"))
+            default = spec.get("default")
+            if default == "":
+                default = None
+            params.append(Param(pname, ty, opt, default, spec.get("description", "") or ""))
+        tools.append(Tool(t["name"], t.get("description", "") or "", params))
+    calls = [Call(a["name"], dict(a.get("arguments") or {})) for a in answers]
+    ex = Example(row["query"], tools, calls, id=f"xlam_{row.get('id', '')}", source="xlam")
+    ex.gold = gold_from_calls(tools, calls)
+    return ex
+
+
+def from_bfcl(q_row: dict, answer_row: dict | None, category: str) -> Example:
+    query = q_row["question"][0][-1]["content"] if q_row["question"] and q_row["question"][0] else ""
+    tools = []
+    for f in q_row["function"]:
+        props = (f.get("parameters") or {}).get("properties", {}) or {}
+        required = set((f.get("parameters") or {}).get("required", []) or [])
+        params = []
+        for pname, spec in props.items():
+            ty, _ = norm_type(spec.get("type"))
+            params.append(Param(pname, ty, pname not in required, spec.get("default"), spec.get("description", "") or "",
+                                spec.get("enum") or None))
+        tools.append(Tool(f["name"], f.get("description", "") or "", params))
+    gold = answer_row["ground_truth"] if answer_row else []
+    calls = []
+    for g in gold:
+        (name, args), = g.items()
+        calls.append(Call(name, {k: v[0] for k, v in args.items() if v and v[0] != ""}))
+    return Example(query, tools, calls, id=q_row["id"], source=f"bfcl_{category}", gold=gold)
+
+
+def gold_from_calls(tools: list[Tool], calls: list[Call]) -> list[dict]:
+    """Build BFCL-style acceptable-value gold from exact xLAM answers."""
+    by_name = {t.name: t for t in tools}
+    gold = []
+    for c in calls:
+        accept: dict[str, list] = {k: [v] for k, v in c.args.items()}
+        tool = by_name.get(c.name)
+        if tool:
+            for p in tool.params:
+                if p.name in accept:
+                    if p.optional and p.default is not None and _same(accept[p.name][0], p.default):
+                        accept[p.name].append("")
+                elif p.optional or p.default is not None:
+                    accept[p.name] = [""] + ([p.default] if p.default is not None else [])
+        gold.append({c.name: accept})
+    return gold
+
+
+def _same(a: Any, b: Any) -> bool:
+    try:
+        return a == b or float(a) == float(b)
+    except (TypeError, ValueError):
+        return str(a) == str(b)
+
+
+MAX_OPTIONS = 16
+
+
+def param_options(p: Param) -> list | None:
+    """Values a v3 option slot picks from: bools -> [True, False], declared enums of <= MAX_OPTIONS, else None."""
+    if p.type == "bool":
+        return [True, False]
+    if p.options and len(p.options) <= MAX_OPTIONS:
+        return list(p.options)
+    return None
+
+
+# ---------- compact signature ----------
+
+def _trim(text: str, words: int) -> str:
+    w = " ".join(str(text).split()).split(" ")
+    return " ".join(w[:words]) if words > 0 else ""
+
+
+def signature_pieces(tool: Tool, desc_words: int = 8) -> list[tuple[str, str, int]]:
+    """Pieces of `name(p:type=default "desc", q?:int) | description`.
+    Each piece is (text, kind, param_index); kind in {'name', 'param', 'text'}.
+    The 'param' piece is exactly the parameter name so its first token can be located."""
+    pieces: list[tuple[str, str, int]] = [(tool.name, "name", -1), ("(", "text", -1)]
+    for k, p in enumerate(tool.params):
+        if k:
+            pieces.append((", ", "text", -1))
+        pieces.append((p.name, "param", k))
+        spec = ("?" if p.optional else "") + ":" + p.type
+        if p.default is not None:
+            spec += "=" + value_to_text(p.default)
+        d = _trim(p.description, desc_words)
+        if d:
+            spec += f' "{d}"'
+        pieces.append((spec, "text", -1))
+    pieces.append((")", "text", -1))
+    d = _trim(tool.description, max(desc_words + 4, 0) if desc_words > 0 else 0)
+    if d:
+        pieces.append((" | " + d, "text", -1))
+    return pieces
+
+
+def signature(tool: Tool, desc_words: int = 8) -> str:
+    return "".join(t for t, _, _ in signature_pieces(tool, desc_words))
+
+
+# ---------- values ----------
+
+def value_to_text(v: Any) -> str:
+    if isinstance(v, str):
+        return v
+    if isinstance(v, bool):
+        return "true" if v else "false"
+    return json.dumps(v, ensure_ascii=False)
+
+
+def text_to_value(s: str, ty: str) -> Any:
+    s = s.strip()
+    if ty == "str":
+        return s
+    if ty == "bool":
+        return s.lower() in ("true", "1", "yes")
+    if ty in ("int", "float"):
+        try:
+            f = float(s)
+            if ty == "int" and f.is_integer():
+                return int(f)
+            return f
+        except ValueError:
+            pass
+    for parse in (json.loads, ast.literal_eval):
+        try:
+            return parse(s)
+        except Exception:
+            continue
+    return s
+
+
+# ---------- OpenAI-style tool (for the LLM baseline) ----------
+
+_JSON_TYPES = {"str": "string", "int": "integer", "float": "number", "bool": "boolean",
+               "list": "array", "dict": "object", "any": "string"}
+
+
+def from_openai_tool(t: dict) -> Tool:
+    """OpenAI-style tool ({"type": "function", "function": {...}} or the bare function dict) -> Tool."""
+    f = t.get("function", t)
+    ps = f.get("parameters") or {}
+    req = set(ps.get("required", []))
+    params = []
+    for name, s in (ps.get("properties") or {}).items():
+        ty, _ = norm_type(s.get("type"))
+        params.append(Param(name, ty, name not in req, s.get("default"), s.get("description", "") or "", s.get("enum")))
+    return Tool(f["name"], f.get("description", "") or "", params)
+
+
+def to_openai_tool(tool: Tool) -> dict:
+    props = {}
+    for p in tool.params:
+        spec = {"type": _JSON_TYPES[p.type], "description": p.description}
+        if p.default is not None:
+            spec["default"] = p.default
+        props[p.name] = spec
+    return {"type": "function", "function": {
+        "name": tool.name, "description": tool.description,
+        "parameters": {"type": "object", "properties": props,
+                       "required": [p.name for p in tool.params if not p.optional]}}}
+
+
+# ---------- (de)serialisation for jsonl ----------
+
+def example_to_dict(ex: Example) -> dict:
+    return {"id": ex.id, "source": ex.source, "query": ex.query, "gold": ex.gold,
+            "tools": [{"name": t.name, "description": t.description,
+                       "params": [vars(p) for p in t.params]} for t in ex.tools],
+            "calls": [vars(c) for c in ex.calls]}
+
+
+def example_from_dict(d: dict) -> Example:
+    tools = [Tool(t["name"], t["description"], [Param(**p) for p in t["params"]]) for t in d["tools"]]
+    return Example(d["query"], tools, [Call(**c) for c in d["calls"]], d["id"], d["source"], d.get("gold", []))

code/actionman/scoring.py

@@ -0,0 +1,87 @@
+"""BFCL-style AST matching (re-implemented, simplified) for predicted calls vs acceptable-value gold.
+
+gold: [{func_name: {arg: [acceptable values]}}]; "" in the acceptable list means the arg may be omitted.
+"""
+from __future__ import annotations
+
+import itertools
+import re
+from typing import Any
+
+from .schema import Call
+
+
+def _std(s: str) -> str:
+    return re.sub(r"[ ,./\-_*^]", "", str(s)).lower().replace("'", '"')
+
+
+def _num(x):
+    if isinstance(x, bool):
+        return None
+    if isinstance(x, (int, float)):
+        return float(x)
+    if isinstance(x, str):
+        try:
+            return float(x)
+        except ValueError:
+            return None
+    return None
+
+
+def value_match(pred: Any, gold: Any) -> bool:
+    if isinstance(gold, bool) or isinstance(pred, bool):
+        if isinstance(pred, str):
+            return pred.strip().lower() == str(gold).lower()
+        return pred == gold
+    gn, pn = _num(gold), _num(pred)
+    if gn is not None and not isinstance(gold, str):
+        return pn is not None and abs(pn - gn) < 1e-6
+    if isinstance(gold, str):
+        return _std(pred) == _std(gold) if not isinstance(pred, (list, dict)) else False
+    if isinstance(gold, (list, tuple)):
+        return isinstance(pred, (list, tuple)) and len(pred) == len(gold) and all(
+            value_match(p, g) for p, g in zip(pred, gold))
+    if isinstance(gold, dict):
+        if not isinstance(pred, dict):
+            return False
+        for k, gv in gold.items():
+            accept = gv if isinstance(gv, list) else [gv]
+            if k not in pred:
+                if "" not in accept:
+                    return False
+                continue
+            if not any(value_match(pred[k], a) for a in accept if a != ""):
+                return False
+        return set(pred) <= set(gold)
+    return pred == gold
+
+
+def call_match(pred: Call, gold_call: dict) -> bool:
+    (name, accept), = gold_call.items()
+    if pred.name != name:
+        return False
+    for k, v in pred.args.items():
+        if k not in accept:
+            return False
+        if not any(value_match(v, a) for a in accept[k] if not (isinstance(a, str) and a == "")):
+            if not ("" in accept[k] and v in ("", None)):
+                return False
+    for k, vals in accept.items():
+        if k not in pred.args and "" not in vals:
+            return False
+    return True
+
+
+def match_calls(pred: list[Call], gold: list[dict]) -> bool:
+    if len(pred) != len(gold):
+        return False
+    if not gold:
+        return True
+    for perm in itertools.permutations(range(len(pred))):
+        if all(call_match(pred[perm[g]], gold[g]) for g in range(len(gold))):
+            return True
+    return False
+
+
+def names_match(pred: list[Call], gold: list[dict]) -> bool:
+    return sorted(c.name for c in pred) == sorted(next(iter(g)) for g in gold)

code/actionman/span.py

@@ -0,0 +1,29 @@
+"""Span pointer decoding for the v3 copy head."""
+from __future__ import annotations
+
+import torch
+
+
+def best_span(start: torch.Tensor, end: torch.Tensor, max_len: int = 32):
+    """Highest p(start) * p(end) with start <= end < start + max_len -> (s, e, prob)."""
+    ps, pe = start.float().softmax(-1), end.float().softmax(-1)
+    Q = ps.shape[0]
+    i = torch.arange(Q, device=ps.device)
+    ok = (i[None, :] >= i[:, None]) & (i[None, :] < i[:, None] + max_len)
+    score = (ps[:, None] * pe[None, :]).masked_fill(~ok, 0)
+    flat = int(score.argmax())
+    return flat // Q, flat % Q, float(score.flatten()[flat])
+
+
+def _cls(ch: str) -> str:
+    return "a" if ch.isalpha() else "d" if ch.isdigit() else ""
+
+
+def snap_word(text: str, s: int, e: int) -> tuple[int, int]:
+    """Widen the char span [s, e) so it never cuts a word or number in half ("US Doll" -> "US Dollars").
+    Only letter-letter and digit-digit edges are extended, so "100" stays "100" inside "100km"."""
+    while s > 0 and s < e and _cls(text[s]) and _cls(text[s - 1]) == _cls(text[s]):
+        s -= 1
+    while e < len(text) and e > s and _cls(text[e - 1]) and _cls(text[e]) == _cls(text[e - 1]):
+        e += 1
+    return s, e

code/actionman/v3feat.py

@@ -0,0 +1,181 @@
+"""v3 sequences: copy targets into the query, Jev-style option blocks, per-item list slots.
+
+Pass B v3 slot block:
+  #1 tool( city= [M][M][M][M] , unit= [M] ( [OPT] metric [OPT] imperial ) , tags= [ [M]x8 | [M]x8 ] )
+Every scalar value / list item has an anchor (its first slot) that carries the copy gate and the span pointer;
+every option param has one [MASK] anchor that points at its [OPT] markers. The v2 builders are untouched.
+"""
+from __future__ import annotations
+
+import torch
+
+from .audit import copy_span, query_offsets
+from .features import BUCKETS, MAX_CALLS, MAX_LEN_B, OVERFLOW, _enc, build_pass_a, value_tokens
+from .schema import Example, param_options, value_to_text
+from .scoring import value_match
+
+ITEM_SLOTS = 8
+MAX_ITEMS = 7
+N_COUNT = MAX_ITEMS + 1          # count classes: 0 = whole list as text (v2 path), 1..7 = item slots
+
+
+def _scalar(v) -> bool:
+    return isinstance(v, (str, int, float)) and not isinstance(v, bool)
+
+
+def list_mode(tok, query: str, v, offs) -> bool:
+    """A list gets item slots if it has 1..7 scalar items, each copyable or at most ITEM_SLOTS tokens."""
+    if not isinstance(v, list) or not 1 <= len(v) <= MAX_ITEMS:
+        return False
+    for it in v:
+        if not _scalar(it):
+            return False
+        if len(value_tokens(tok, it)) > ITEM_SLOTS and copy_span(tok, query, value_to_text(it), offs) is None:
+            return False
+    return True
+
+
+def option_index(p, v) -> int | None:
+    opts = param_options(p)
+    if opts is None:
+        return None
+    for m, o in enumerate(opts):
+        if value_match(v, o) or value_to_text(v).strip().lower() == value_to_text(o).strip().lower():
+            return m
+    return None
+
+
+def build_pass_a_v3(tok, ex: Example) -> dict | None:
+    """Pass A plus query token positions/char offsets and count labels for declared-list params."""
+    a = build_pass_a(tok, ex)
+    if a is None:
+        return None
+    offs = query_offsets(tok, ex.query)
+    q_end = a["call_pos"][0]                       # query tokens sit right before the [CALL] markers
+    a["q_pos"] = list(range(q_end - len(offs), q_end))
+    a["q_offs"] = offs
+    count_items = []
+    for (i, j, k, b) in a["len_items"]:
+        p = ex.tools[j].params[k]
+        if p.type == "list" and b != 0:
+            v = ex.calls[i].args[p.name]
+            count_items.append((i, j, k, len(v) if list_mode(tok, ex.query, v, offs) else 0))
+    a["count_items"] = count_items
+    return a
+
+
+def gold_plan_v3(tok, ex: Example, a: dict) -> list[tuple[int, dict[int, tuple[str, int]]]]:
+    """Per gold call: (tool_j, {param_k: (mode, n)}) with mode in opt / list (n items) / text (n = bucket)."""
+    plan = []
+    for i in range(min(len(ex.calls), MAX_CALLS)):
+        j, modes = a["dec_labels"][i], {}
+        for (ci, _, k, b) in a["len_items"]:
+            if ci != i or b == 0:
+                continue
+            p = ex.tools[j].params[k]
+            v = ex.calls[i].args[p.name]
+            if option_index(p, v) is not None:
+                modes[k] = ("opt", 0)
+            elif p.type == "list" and list_mode(tok, ex.query, v, a["q_offs"]):
+                modes[k] = ("list", len(v))
+            else:
+                modes[k] = ("text", b)
+        plan.append((j, modes))
+    return plan
+
+
+def build_pass_b_v3(tok, ex: Example, a: dict, plan, with_targets: bool, max_len: int = MAX_LEN_B,
+                    keep_context: bool = False, markers: bool = True) -> dict | None:
+    ids = (a["input_ids"] if keep_context else a["input_ids"][:-1]) + [tok.am["[ARGS]"]]
+    ctx_len = len(a["input_ids"]) if keep_context else len(ids) - 1
+    out = {"slot_pos": [], "slot_index": [], "anchors": [], "opt_slots": [], "ctx_len": ctx_len, "q_pos": a["q_pos"]}
+    if with_targets:
+        out["slot_targets"], out["copy_targets"] = [], []
+
+    def add_value(n, value, key):
+        pos = list(range(len(ids), len(ids) + n))
+        ids.extend([tok.mask_token_id] * n)
+        out["slot_pos"].append(pos)
+        out["slot_index"].append(key)
+        out["anchors"].append(pos[0])
+        if with_targets:
+            v = value_tokens(tok, value)[:n]
+            out["slot_targets"].append(v + [tok.pad_token_id] * (n - len(v)))
+            cs = copy_span(tok, ex.query, value_to_text(value), a["q_offs"]) if _scalar(value) else None
+            out["copy_targets"].append((a["q_pos"][cs[0]], a["q_pos"][cs[1]]) if cs else None)
+
+    for i, (j, modes) in enumerate(plan):
+        tool = ex.tools[j]
+        ids.extend(_enc(tok, (f" #{i + 1} " if markers else "") + tool.name + "("))
+        first = True
+        for k, p in enumerate(tool.params):
+            if k not in modes:
+                continue
+            mode, n = modes[k]
+            if mode == "text" and n in (0, OVERFLOW):
+                continue
+            ids.extend(_enc(tok, ("" if first else ", ") + p.name + "="))
+            first = False
+            v = ex.calls[i].args[p.name] if with_targets else None
+            if mode == "opt":
+                anchor = len(ids)
+                ids.append(tok.mask_token_id)
+                ids.extend(_enc(tok, " ("))
+                opos = []
+                for o in param_options(p):
+                    opos.append(len(ids))
+                    ids.append(tok.am["[OPT]"])
+                    ids.extend(_enc(tok, " " + value_to_text(o)))
+                ids.extend(_enc(tok, " )"))
+                tgt = option_index(p, v) if with_targets else None
+                out["opt_slots"].append((i, k, anchor, opos, -100 if tgt is None else tgt))
+            elif mode == "list":
+                ids.extend(_enc(tok, " ["))
+                for it in range(n):
+                    if it:
+                        ids.extend(_enc(tok, " |"))
+                    add_value(ITEM_SLOTS, v[it] if with_targets else None, (i, k, it))
+                ids.extend(_enc(tok, " ]"))
+            else:
+                add_value(BUCKETS[n], v, (i, k, None))
+        ids.extend(_enc(tok, ")"))
+    ids.append(tok.sep_token_id)
+    if len(ids) > max_len:
+        return None
+    out["input_ids"] = ids
+    return out
+
+
+# ---------- batching ----------
+
+def v3_extras(A: list[dict], B_rows: list[tuple[int, dict]]) -> dict:
+    """Tensors for the v3 heads. A: pass-A dicts (row = index); B_rows: (row in the pass-B batch, pass-B dict)."""
+    ab, ap, gy, ss, se, ob, oa, oy, opos = [], [], [], [], [], [], [], [], []
+    Q = max([len(b["q_pos"]) for _, b in B_rows] + [1])
+    n_rows = max([r for r, _ in B_rows] + [-1]) + 1
+    qpos = torch.full((max(n_rows, 1), Q), -1, dtype=torch.long)
+    for r, b in B_rows:
+        qpos[r, : len(b["q_pos"])] = torch.tensor(b["q_pos"], dtype=torch.long)
+        q0 = b["q_pos"][0] if b["q_pos"] else 0
+        clean = set(b.get("_clean", b["anchors"]))     # set by data._mask_slots for training batches
+        for anchor, ct in zip(b["anchors"], b.get("copy_targets", [None] * len(b["anchors"]))):
+            if anchor not in clean:
+                continue
+            ab.append(r); ap.append(anchor)
+            gy.append(1.0 if ct else 0.0)
+            ss.append(ct[0] - q0 if ct else -100); se.append(ct[1] - q0 if ct else -100)
+        for (_, _, anchor, opts, tgt) in b["opt_slots"]:
+            ob.append(r); oa.append(anchor); oy.append(tgt); opos.append(opts)
+    M = max([len(o) for o in opos] + [1])
+    opt_pos = torch.full((len(opos), M), -1, dtype=torch.long)
+    for n, o in enumerate(opos):
+        opt_pos[n, : len(o)] = torch.tensor(o, dtype=torch.long)
+    cidx, cy = [], []
+    for i, a in enumerate(A):
+        for (ci, j, k, c) in a.get("count_items", []):
+            cidx.append((i, a["call_pos"][ci], a["arg_pos"][(j, k)]))
+            cy.append(c)
+    L = lambda x, dt=torch.long: torch.tensor(x, dtype=dt)  # noqa: E731
+    return {"anchor_b": L(ab), "anchor_p": L(ap), "gate_y": L(gy, torch.float), "span_s": L(ss), "span_e": L(se),
+            "qpos": qpos, "opt_b": L(ob), "opt_anchor": L(oa), "opt_pos": opt_pos, "opt_y": L(oy),
+            "count_idx": L(cidx).reshape(-1, 3), "count_y": L(cy)}

config.json

@@ -0,0 +1,107 @@
+{
+ "version": "0.1",
+ "base": "answerdotai/ModernBERT-large",
+ "flags": {
+  "cached": true,
+  "markers": true,
+  "v3": true
+ },
+ "encoder": {
+  "transformers_version": "5.19.0",
+  "architectures": [
+   "ModernBertForMaskedLM"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": "float32",
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+   "0": "LABEL_0",
+   "1": "LABEL_1"
+  },
+  "label2id": {
+   "LABEL_0": 0,
+   "LABEL_1": 1
+  },
+  "problem_type": null,
+  "vocab_size": 50432,
+  "hidden_size": 1024,
+  "intermediate_size": 2624,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "hidden_activation": "gelu",
+  "max_position_embeddings": 8192,
+  "initializer_range": 0.02,
+  "initializer_cutoff_factor": 2.0,
+  "norm_eps": 1e-05,
+  "norm_bias": false,
+  "pad_token_id": 50283,
+  "eos_token_id": 50282,
+  "bos_token_id": 50281,
+  "cls_token_id": 50281,
+  "sep_token_id": 50282,
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "layer_types": [
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention",
+   "sliding_attention",
+   "sliding_attention",
+   "full_attention"
+  ],
+  "rope_parameters": {
+   "full_attention": {
+    "rope_theta": 160000.0,
+    "rope_type": "default"
+   },
+   "sliding_attention": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+   }
+  },
+  "local_attention": 128,
+  "embedding_dropout": 0.0,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "decoder_bias": true,
+  "classifier_pooling": "mean",
+  "classifier_dropout": 0.0,
+  "classifier_bias": false,
+  "classifier_activation": "gelu",
+  "deterministic_flash_attn": false,
+  "sparse_prediction": false,
+  "sparse_pred_ignore_index": -100,
+  "tie_word_embeddings": true,
+  "_name_or_path": "",
+  "global_attn_every_n_layers": 3,
+  "gradient_checkpointing": false,
+  "layer_norm_eps": 1e-05,
+  "model_type": "modernbert",
+  "position_embedding_type": "absolute",
+  "output_attentions": false
+ }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2e10ab52163263397dfe356244cdbed2487488962695604879eb3979efb993d
+size 924614024

tokenizer_config.json

@@ -0,0 +1,24 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": [
+    "[T]",
+    "[CALL]",
+    "[ARGS]",
+    "[USER]",
+    "[OPT]"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "[MASK]",
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "[UNK]"
+}