caudate-cli 0.1.0__py3-none-any.whl

nn/data.py

@@ -0,0 +1,1635 @@
+"""Dataset + replay buffer for the cognitive controller.
+
+Two data sources:
+
+  1. **Disk corpus** — every saved Session under `data/sessions/*.json`.
+     Each tool call inside a session becomes one (state, action, reward)
+     example. Reward defaults to 0.5 unless the session has an explicit
+     reflection score attached.
+
+  2. **Live replay buffer** — circular in-memory buffer fed by the agent
+     during use. Lets the controller train online from current sessions
+     before they're saved to disk.
+
+ToolVocab maps the dynamic set of registered tools to stable integer
+ids that survive across runs. The vocab grows append-only — adding a
+new tool doesn't invalidate older checkpoints; it just adds a row to
+the embedding table at training time.
+"""
+
+from __future__ import annotations
+
+import collections
+import json
+import logging
+import random
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Iterator
+
+import torch
+
+from nn.config import NNConfig
+from nn.format import ChatMessage, Conversation, ToolCall, ToolDef
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------
+# Tool vocab
+# ---------------------------------------------------------------------
+
+
+class ToolVocab:
+    """Bidirectional id ↔ name map for tool tokens.
+
+    Four special ids are reserved upfront:
+      <pad>      — padding for batched sequences
+      <unk>      — placeholder for missing/unmapped labels at inference
+      <bos>      — start-of-sequence
+      <no_tool>  — explicit "no tool was called" — a real action class
+                   (distinct from <unk>: a turn where the assistant
+                   correctly chose to just answer without any tool)
+
+    Keeping <no_tool> separate from <unk> matters for training: ~40% of
+    real conversation turns don't call a tool, and that's a valid
+    decision the model should learn to *predict*, not a degenerate
+    fallback class.
+    """
+
+    SPECIAL = {
+        "<pad>": 0,
+        "<unk>": 1,
+        "<bos>": 2,
+        "<no_tool>": 3,
+    }
+
+    def __init__(self):
+        self._id2name: dict[int, str] = {v: k for k, v in self.SPECIAL.items()}
+        self._name2id: dict[str, int] = dict(self.SPECIAL)
+        self._next_id = max(self.SPECIAL.values()) + 1
+
+    def add(self, name: str) -> int:
+        if name in self._name2id:
+            return self._name2id[name]
+        idx = self._next_id
+        self._next_id += 1
+        self._name2id[name] = idx
+        self._id2name[idx] = name
+        return idx
+
+    def get(self, name: str) -> int:
+        return self._name2id.get(name, self.SPECIAL["<unk>"])
+
+    def name(self, idx: int) -> str:
+        return self._id2name.get(idx, "<unk>")
+
+    def __len__(self) -> int:
+        return self._next_id
+
+    def to_dict(self) -> dict[str, int]:
+        return dict(self._name2id)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, int]) -> "ToolVocab":
+        v = cls()
+        for name, idx in data.items():
+            if name in v.SPECIAL:
+                continue
+            v._name2id[name] = idx
+            v._id2name[idx] = name
+            v._next_id = max(v._next_id, idx + 1)
+        return v
+
+
+class SourceVocab:
+    """Bidirectional id ↔ name map for teacher-model sources.
+
+    Phase 2 of CAUDATE_EVOLUTION.md: every ConversationSample carries a
+    `model_source` string ("anthropic/claude-opus-4-7", "ollama/...",
+    "<unknown>", etc.). The model conditions on that source via a
+    learned embedding lookup. This vocab maps the names to stable ids,
+    capped at `cfg.source_vocab_size` so the embedding table stays
+    finite. Overflow (more than the cap) folds back to <unknown>.
+
+    Reserved slot 0 = <unknown> = the zero-bias baseline. Legacy
+    untagged samples and any model we haven't seen before share this
+    slot.
+    """
+
+    SPECIAL = {"<unknown>": 0}
+    DEFAULT_CAP = 16
+
+    def __init__(self, cap: int = DEFAULT_CAP):
+        self.cap = max(1, int(cap))
+        self._id2name: dict[int, str] = {v: k for k, v in self.SPECIAL.items()}
+        self._name2id: dict[str, int] = dict(self.SPECIAL)
+        self._next_id = max(self.SPECIAL.values()) + 1
+
+    def add(self, name: str) -> int:
+        if not isinstance(name, str) or not name:
+            return self.SPECIAL["<unknown>"]
+        if name in self._name2id:
+            return self._name2id[name]
+        if self._next_id >= self.cap:
+            # Overflow: silently fold to <unknown>. Better than
+            # blowing up; the embedding table can't grow at runtime.
+            return self.SPECIAL["<unknown>"]
+        idx = self._next_id
+        self._next_id += 1
+        self._name2id[name] = idx
+        self._id2name[idx] = name
+        return idx
+
+    def get(self, name: str) -> int:
+        return self._name2id.get(name, self.SPECIAL["<unknown>"])
+
+    def name(self, idx: int) -> str:
+        return self._id2name.get(idx, "<unknown>")
+
+    def __len__(self) -> int:
+        return self._next_id
+
+    def to_dict(self) -> dict[str, int]:
+        return dict(self._name2id)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, int], cap: int = DEFAULT_CAP) -> "SourceVocab":
+        v = cls(cap=cap)
+        for name, idx in data.items():
+            if name in v.SPECIAL:
+                continue
+            if int(idx) >= v.cap:
+                continue   # truncate any persisted ids beyond the new cap
+            v._name2id[name] = int(idx)
+            v._id2name[int(idx)] = name
+            v._next_id = max(v._next_id, int(idx) + 1)
+        return v
+
+
+# ---------------------------------------------------------------------
+# Sample structure — ConversationSample is THE training row.
+# ---------------------------------------------------------------------
+# Standard chat-tool-call schema (OpenAI shape) — same format used by
+# every public function-calling dataset (xLAM, ToolBench, ToolAlpaca,
+# Gorilla, API-Bank). Cognos sessions, external HuggingFace datasets,
+# and the agent's live replay buffer all produce this single type, so
+# the training path doesn't need to know where samples came from.
+#
+# Layout:
+#   - conversation:  list[ChatMessage]   ← the role-tagged turn list
+#   - tools:         list[ToolDef]       ← what was available to call
+#   - target_tool:   str                 ← name of the tool actually called
+#   - target_arguments: str              ← JSON-encoded args (reserved for
+#                                           a future generative arg head)
+#   - target_tool_call_index: int        ← which tool_call within the
+#                                           target assistant message
+#   - Cognos-specific signals (mood, model_source, tier/think/value, the
+#     14 optional D-heads) layer on top. External data leaves them at
+#     their default values; Cognos sessions populate them.
+
+
+@dataclass
+class ConversationSample:
+    """One training example in standard chat-tool-call format.
+
+    The conversation + tools fields together form one row that any
+    OpenAI-style tool-use dataset can produce. Cognos signals layer
+    on top as optional fields — they stay None for external data, so
+    the trainer's optional_target machinery skips those heads cleanly.
+    """
+
+    # ── Standard-format payload ──────────────────────────────────────
+    conversation: list[ChatMessage] = field(default_factory=list)
+    tools: list[ToolDef] = field(default_factory=list)
+
+    # ── Targets ──────────────────────────────────────────────────────
+    # The tool name the assistant called at the target turn. Empty
+    # string or "<no_tool>" for sample-points where the assistant
+    # answered directly. Both forms are legal — the collate path
+    # normalises via the tool vocab.
+    target_tool: str = "<no_tool>"
+    target_arguments: str = ""        # JSON; reserved for Phase 2
+    target_tool_call_index: int = 0    # which tool_call within the turn
+
+    # ── Cognos signals (optional) ────────────────────────────────────
+    mood: list[float] = field(default_factory=lambda: [0.5, 0.5, 0.5, 0.5])
+    image_paths: list[str] = field(default_factory=list)
+    model_source: str = "<unknown>"
+    surprise: float = 0.5
+    target_tier: int = 0
+    target_think: float = 0.0
+    target_value: float = 0.0
+    # Optional D-heads + Tier-1/2/4b heads (all may be None)
+    target_memory_write: float | None = None
+    target_cache_hit: float | None = None
+    target_permission: float | None = None
+    target_refusal: float | None = None
+    target_code_response: float | None = None
+    target_stall: float | None = None
+    target_difficulty: int | None = None
+    target_stop_iter: float | None = None
+    target_compaction: float | None = None
+    target_latency_s: float | None = None
+    target_token_budget: float | None = None
+    target_mood_pred: list[float] | None = None
+    target_subagent_spawn: float | None = None
+    target_reward_model: float | None = None
+    target_feature_success: float | None = None
+
+    # ── Construction helpers ─────────────────────────────────────────
+
+    @classmethod
+    def from_conversation(
+        cls,
+        conv: Conversation,
+        prefix_len: int,
+        tool_call: ToolCall,
+        tool_call_index: int = 0,
+        **cognos_signals: Any,
+    ) -> "ConversationSample":
+        """Build one sample by taking the first `prefix_len` messages
+        as context and the given tool_call as the target.
+
+        `prefix_len` is the number of messages BEFORE the assistant
+        message that emitted this tool_call. The prefix is what the
+        model sees; the target_tool is what it should predict.
+        Cognos signals (mood, model_source, etc.) are passed via
+        kwargs so the loader stays in control of which ones it sets.
+        """
+        prefix = list(conv.messages[:prefix_len])
+        return cls(
+            conversation=prefix,
+            tools=list(conv.tools),
+            target_tool=tool_call.name or "<no_tool>",
+            target_arguments=tool_call.arguments or "",
+            target_tool_call_index=tool_call_index,
+            **cognos_signals,
+        )
+
+
+# ---------------------------------------------------------------------
+# Conversion helpers used by collate & every loader
+# ---------------------------------------------------------------------
+
+
+def conversation_to_strings(conv: list[ChatMessage]) -> list[str]:
+    """Convert a list[ChatMessage] into the role-prefixed string list
+    the StateEncoder consumes.
+
+    Tool calls inside an assistant message are rendered as
+    ``"assistant: <content> calls: tool_name(args), ...".`` Tool result
+    messages (role="tool") are rendered as ``"tool {name}: <content>".``
+    Empty content + no calls is dropped (a no-op turn carries no signal).
+
+    The encoder operates on text; this helper is the boundary that
+    keeps the chat-schema clean upstream and the encoder API
+    unchanged downstream.
+    """
+    out: list[str] = []
+    for m in conv:
+        if m.role == "tool":
+            label = m.name or "tool"
+            content = (m.content or "").strip()
+            if content:
+                out.append(f"tool {label}: {content[:400]}")
+            continue
+        if m.tool_calls:
+            call_strs = []
+            for tc in m.tool_calls:
+                arg_preview = (tc.arguments or "").replace("\n", " ")[:60]
+                call_strs.append(f"{tc.name}({arg_preview})")
+            content = (m.content or "").strip()
+            text = (content + " " if content else "") + "calls: " + ", ".join(call_strs)
+            out.append(f"{m.role}: {text[:400]}")
+            continue
+        content = (m.content or "").strip()
+        if content:
+            out.append(f"{m.role}: {content[:400]}")
+    return out
+
+
+def conversation_tool_history(conv: list[ChatMessage]) -> list[str]:
+    """Extract the sequence of tool names called in this conversation."""
+    history: list[str] = []
+    for m in conv:
+        if m.role == "assistant":
+            for tc in m.tool_calls:
+                if tc.name:
+                    history.append(tc.name)
+    return history
+
+
+# ---------------------------------------------------------------------
+# Cognos-session loader — reads data/sessions/*.json into ConversationSamples
+# ---------------------------------------------------------------------
+
+
+def load_corpus_from_sessions(
+    sessions_dir: Path = Path("data/sessions"),
+) -> list[ConversationSample]:
+    """Convert each saved Cognos session into ConversationSample(s).
+
+    Sessions are stored as OpenAI-shape JSON. We emit one
+    ConversationSample per assistant tool call, with the prefix-of-
+    messages-before-the-call as context. Reward semantics:
+    metadata.reflection_scores[call_id] → target_value, else 0.5.
+
+    Sessions that recorded their tool registry under metadata.tools
+    have that list carried through; absent tools means empty (the
+    open-vocab head still trains, just without negative candidates).
+    """
+    out: list[ConversationSample] = []
+    if not sessions_dir.exists():
+        return out
+
+    for path in sorted(sessions_dir.glob("*.json")):
+        try:
+            data = json.loads(path.read_text())
+        except Exception:
+            continue
+        raw_msgs = data.get("messages") or []
+        scores = (data.get("metadata") or {}).get("reflection_scores") or {}
+        raw_tools = (data.get("metadata") or {}).get("tools") \
+            or data.get("tools") or []
+        try:
+            conv = Conversation.from_dict(
+                {"messages": raw_msgs, "tools": raw_tools},
+            )
+        except Exception:
+            continue
+
+        for msg_idx, msg in enumerate(conv.messages):
+            if msg.role != "assistant" or not msg.tool_calls:
+                continue
+            for call_idx, tc in enumerate(msg.tool_calls):
+                if not tc.name:
+                    continue
+                score = float(scores.get(tc.id, 0.5)) if tc.id else 0.5
+                sample = ConversationSample.from_conversation(
+                    conv=conv,
+                    prefix_len=msg_idx,
+                    tool_call=tc,
+                    tool_call_index=call_idx,
+                    model_source=(data.get("metadata") or {}).get(
+                        "model_source", "<unknown>",
+                    ),
+                    target_tier=int(
+                        (data.get("metadata") or {}).get("last_tier", 0),
+                    ),
+                    target_think=float(
+                        (data.get("metadata") or {}).get("thinking_used", 0.0),
+                    ),
+                    target_value=score,
+                )
+                out.append(sample)
+    return out
+
+
+# ---------------------------------------------------------------------
+# OpenAI-format JSONL loader (external datasets: xLAM, ToolBench, etc.)
+# ---------------------------------------------------------------------
+
+
+def load_corpus_from_oai_jsonl(
+    path: Path,
+    model_source: str = "<external>",
+    max_rows: int | None = None,
+) -> list[ConversationSample]:
+    """Load a JSONL file in standard chat-tool-call format.
+
+    One line per conversation. Each conversation may contain multiple
+    assistant tool calls; we emit one ConversationSample per tool call,
+    with the prefix-of-messages-before-the-call as context.
+
+    The standard format works for: xLAM, ToolBench, ToolAlpaca, Gorilla,
+    API-Bank, and any OpenAI-fine-tuning-shaped corpus. Some datasets
+    use ``conversations`` instead of ``messages`` for the turn list —
+    ``Conversation.from_dict`` handles both.
+
+    Malformed rows are skipped, not raised — external corpora are noisy
+    and one bad row must not abort the whole training pass.
+
+    Args:
+        path: JSONL file path.
+        model_source: tagged on every emitted sample so the
+            source-embedding (Phase 2 of CAUDATE_EVOLUTION) can bias
+            predictions for external-data turns.
+        max_rows: cap for quick smoke runs; None = read everything.
+
+    Returns:
+        list of ConversationSample, one per assistant tool call.
+    """
+    out: list[ConversationSample] = []
+    if not path.exists():
+        logger.warning(f"oai jsonl not found: {path}")
+        return out
+
+    n_rows = 0
+    n_calls = 0
+    with path.open(encoding="utf-8") as f:
+        for line in f:
+            if max_rows is not None and n_rows >= max_rows:
+                break
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                row = json.loads(line)
+            except Exception:
+                continue
+            try:
+                conv = Conversation.from_dict(row)
+            except Exception as e:
+                logger.debug(f"skipping malformed oai row: {e}")
+                continue
+            if not conv.messages:
+                continue
+            n_rows += 1
+
+            # Row-level target_tool: the cognos-toolbox amplified/replay
+            # rows are "predict the next tool given user turn + tool
+            # catalog" — no assistant message exists yet. Emit one
+            # sample with the row-level target and the full message
+            # list as context.
+            row_target = row.get("target_tool")
+            if row_target is not None:
+                row_meta = row.get("meta") or {}
+                row_source = (
+                    str(row_meta.get("teacher_model"))
+                    if isinstance(row_meta, dict) and row_meta.get("teacher_model")
+                    else model_source
+                )
+                out.append(ConversationSample(
+                    conversation=list(conv.messages),
+                    tools=list(conv.tools),
+                    target_tool=str(row_target) or "<no_tool>",
+                    model_source=row_source,
+                ))
+                n_calls += 1
+                continue
+
+            # Standard OAI shape: emit one sample per tool_call within
+            # each assistant message. The prefix is everything *before*
+            # that assistant message — that's the state the model sees
+            # when deciding which tool to call.
+            for msg_idx, msg in enumerate(conv.messages):
+                if msg.role != "assistant" or not msg.tool_calls:
+                    continue
+                for call_idx, tc in enumerate(msg.tool_calls):
+                    if not tc.name:
+                        continue
+                    sample = ConversationSample.from_conversation(
+                        conv=conv,
+                        prefix_len=msg_idx,
+                        tool_call=tc,
+                        tool_call_index=call_idx,
+                        model_source=model_source,
+                    )
+                    out.append(sample)
+                    n_calls += 1
+
+    logger.info(
+        f"loaded {n_rows} rows / {n_calls} tool-call samples from {path}"
+    )
+    return out
+
+
+# ---------------------------------------------------------------------
+# HuggingFace dataset loader — direct path from HF Hub to ConversationSample
+# ---------------------------------------------------------------------
+#
+# HF function-calling datasets don't share one schema. Common shapes:
+#
+#   OpenAI chat:   row = {"messages": [...], "tools": [...]}
+#   xLAM-style:    row = {"query": "...", "tools": "[...]",
+#                         "answers": "[{...tool_call...}]"}     (JSON strings)
+#   Conversations: row = {"conversations": [...same as messages...],
+#                         "tools": [...]}
+#
+# `load_corpus_from_hf` auto-detects which one applies per dataset. For
+# anything off the beaten path, pass a `row_to_conversation` callable
+# that returns a Conversation given a single row — bypasses detection
+# entirely.
+
+
+def _detect_hf_shape(row: dict[str, Any]) -> str:
+    """Return 'openai' | 'block_content' | 'conversations' | 'sharegpt'
+    | 'xlam' | 'unknown'.
+
+    The ``block_content`` shape is used by llamafactory's
+    reason-tool-use-demo-style corpora: ``messages[].content`` is a
+    list of ``{type, value}`` blocks (``text``/``reasoning``/
+    ``tool_call``) and ``tools`` is a JSON-encoded string rather than
+    a list. Detected ahead of plain ``openai`` because the row's
+    ``messages[0]`` still has a ``role`` key — without the early
+    check we'd silently drop every tool_call.
+    """
+    if isinstance(row.get("messages"), list) and row["messages"] \
+            and isinstance(row["messages"][0], dict) \
+            and "role" in row["messages"][0]:
+        # Distinguish OAI-shape (content: str|None) from block-typed
+        # content (content: list of {type, value} dicts).
+        first_content = row["messages"][0].get("content")
+        if (isinstance(first_content, list) and first_content
+                and isinstance(first_content[0], dict)
+                and "value" in first_content[0]):
+            return "block_content"
+        return "openai"
+    if isinstance(row.get("conversations"), list) and row["conversations"] \
+            and isinstance(row["conversations"][0], dict):
+        first = row["conversations"][0]
+        if "role" in first:
+            return "conversations"
+        if "from" in first and "value" in first:
+            return "sharegpt"
+    if "query" in row and "answers" in row:
+        return "xlam"
+    if isinstance(row.get("system"), str) and isinstance(row.get("chat"), str):
+        return "glaive"
+    return "unknown"
+
+
+def _parse_block_content_row(row: dict[str, Any]) -> Conversation:
+    """Parse a llamafactory-style row where each message's ``content``
+    is a list of typed blocks.
+
+    Block-type mapping:
+      - ``text`` → appended to message content
+      - ``reasoning`` → appended to message content (the trunk sees
+        it as additional context; semantically a "chain of thought"
+        rendered into the same text channel)
+      - ``tool_call`` → parsed as JSON, becomes a ToolCall on this
+        assistant message
+
+    The row's ``tools`` field is a JSON string; we decode it before
+    handing to Conversation.from_dict. When parsing fails or the
+    field is empty, the tools list is left empty.
+    """
+    raw_msgs: list[dict[str, Any]] = []
+    for m in row.get("messages") or []:
+        if not isinstance(m, dict):
+            continue
+        role = str(m.get("role") or "user")
+        content_parts: list[str] = []
+        tool_calls: list[dict[str, Any]] = []
+        blocks = m.get("content")
+        if isinstance(blocks, list):
+            for blk in blocks:
+                if not isinstance(blk, dict):
+                    continue
+                btype = str(blk.get("type") or "").lower()
+                value = blk.get("value") or blk.get("text") or ""
+                if btype == "tool_call":
+                    # value is a JSON string like
+                    # '{"name": "...", "arguments": {...}}'
+                    try:
+                        parsed = json.loads(value) if isinstance(value, str) else value
+                    except Exception:
+                        parsed = None
+                    if isinstance(parsed, dict) and parsed.get("name"):
+                        args = parsed.get("arguments")
+                        if isinstance(args, dict):
+                            args = json.dumps(args)
+                        elif args is None:
+                            args = ""
+                        tool_calls.append({
+                            "type": "function",
+                            "function": {
+                                "name": str(parsed["name"]),
+                                "arguments": str(args),
+                            },
+                        })
+                elif btype in ("text", "reasoning"):
+                    if isinstance(value, str) and value:
+                        content_parts.append(value)
+                else:
+                    # Unknown block types — preserve as text rather
+                    # than drop. Better to feed noisy context than
+                    # lose a labelled message.
+                    if isinstance(value, str) and value:
+                        content_parts.append(value)
+        elif isinstance(blocks, str):
+            content_parts.append(blocks)
+
+        msg_dict: dict[str, Any] = {
+            "role": role,
+            "content": " ".join(content_parts).strip(),
+        }
+        if tool_calls:
+            msg_dict["tool_calls"] = tool_calls
+        raw_msgs.append(msg_dict)
+
+    # tools is a JSON-encoded string in this shape; decode.
+    raw_tools_field = row.get("tools")
+    raw_tools: list[dict[str, Any]] = []
+    if isinstance(raw_tools_field, str) and raw_tools_field.strip():
+        try:
+            decoded = json.loads(raw_tools_field)
+            if isinstance(decoded, list):
+                raw_tools = [t for t in decoded if isinstance(t, dict)]
+        except Exception:
+            pass
+    elif isinstance(raw_tools_field, list):
+        raw_tools = [t for t in raw_tools_field if isinstance(t, dict)]
+
+    return Conversation.from_dict({"messages": raw_msgs, "tools": raw_tools})
+
+
+# Map ShareGPT speaker tags to standard role names. ShareGPT uses
+# `human`/`gpt` historically; modern variants also include `system`,
+# `tool`, and `function`. Unknown tags fall back to "user" rather
+# than silently dropping the turn — better to keep noisy context than
+# lose a labelled message.
+_SHAREGPT_ROLE_MAP: dict[str, str] = {
+    "human": "user",
+    "user": "user",
+    "gpt": "assistant",
+    "chatgpt": "assistant",
+    "bing": "assistant",
+    "assistant": "assistant",
+    "system": "system",
+    "tool": "tool",
+    "function": "tool",
+    "observation": "tool",
+}
+
+
+def _parse_sharegpt_row(row: dict[str, Any]) -> Conversation:
+    """Map a ShareGPT-format row (``conversations`` of ``{from, value}``)
+    to a standard ``Conversation``.
+
+    Used by datasets like Kimi-K2.5-Reasoning, Vicuna, WizardLM,
+    Alpaca-fc-* etc. These are typically reasoning corpora without
+    tool calls — the loader still emits one ConversationSample per
+    assistant turn (target_tool=``<no_tool>``) when
+    ``emit_no_tool_turns=True`` so the trunk gets training signal.
+    """
+    msgs: list[dict[str, Any]] = []
+    for m in row.get("conversations") or []:
+        if not isinstance(m, dict):
+            continue
+        role = _SHAREGPT_ROLE_MAP.get(str(m.get("from") or "").lower(), "user")
+        msgs.append({"role": role, "content": m.get("value") or ""})
+    return Conversation.from_dict({"messages": msgs, "tools": []})
+
+
+def _parse_xlam_row(row: dict[str, Any]) -> Conversation:
+    """xLAM rows pack tools and the assistant's tool calls into JSON
+    strings on separate columns. Rebuild a 2-turn conversation:
+
+        user: <query> ─► assistant: <tool_calls=...>
+
+    so the standard pipeline sees one assistant tool-call turn. We
+    deliberately drop tool-result content (xLAM doesn't simulate them)
+    — the model only learns to PREDICT the call, which is what
+    Caudate's advisor head does anyway.
+    """
+    def _maybe_json(v):
+        if isinstance(v, str):
+            try:
+                return json.loads(v)
+            except Exception:
+                return []
+        return v or []
+
+    tools_raw = _maybe_json(row.get("tools"))
+    answers = _maybe_json(row.get("answers"))
+    if not isinstance(answers, list):
+        answers = []
+    if not isinstance(tools_raw, list):
+        tools_raw = []
+
+    tool_calls_raw: list[dict[str, Any]] = []
+    for a in answers:
+        if not isinstance(a, dict):
+            continue
+        # xLAM uses {"name": ..., "arguments": {...}}
+        tool_calls_raw.append({
+            "type": "function",
+            "function": {
+                "name": a.get("name") or a.get("tool") or "",
+                "arguments": a.get("arguments") or a.get("args") or {},
+            },
+        })
+
+    msgs: list[dict[str, Any]] = [
+        {"role": "user", "content": str(row.get("query") or "")},
+        {"role": "assistant", "content": "", "tool_calls": tool_calls_raw},
+    ]
+    # xLAM's tool defs already look like {"name", "description", "parameters"}
+    # — wrap to OpenAI shape and let Conversation.from_dict normalise.
+    tools_norm = [
+        {"type": "function", "function": t if isinstance(t, dict) else {}}
+        for t in tools_raw
+    ]
+    return Conversation.from_dict({"messages": msgs, "tools": tools_norm})
+
+
+def load_corpus_from_hf(
+    dataset_name: str,
+    split: str = "train",
+    model_source: str | None = None,
+    max_rows: int | None = None,
+    row_to_conversation: "Any | None" = None,
+    streaming: bool = True,
+    emit_no_tool_turns: bool = True,
+    **load_dataset_kwargs: Any,
+) -> list[ConversationSample]:
+    """Stream a HuggingFace dataset into ConversationSamples.
+
+    Auto-detects the row layout. Supported out-of-the-box:
+
+      - **OpenAI chat** — ``messages``/``tools`` columns, dicts with
+        ``role``. Examples: anything saved via the OpenAI fine-tuning
+        format or Hermes-Function-Calling.
+      - **Conversations alias** — same as OpenAI but the column is
+        called ``conversations`` (ToolBench-style).
+      - **xLAM** — ``query``/``tools``/``answers`` columns where
+        ``tools`` and ``answers`` are JSON-encoded strings. Examples:
+        Salesforce/xlam-function-calling-60k.
+
+    For any other shape, pass ``row_to_conversation=fn`` where ``fn``
+    takes a row dict and returns a ``Conversation``. The auto-detector
+    is bypassed in that case.
+
+    Args:
+        dataset_name: HF Hub identifier, e.g.
+            ``"Salesforce/xlam-function-calling-60k"``.
+        split: split name (``"train"``, ``"validation"``, etc.)
+        model_source: tag set on every sample's ``model_source``.
+            Defaults to ``"hf:<dataset_name>"``.
+        max_rows: cap for smoke runs. None = read everything.
+        row_to_conversation: optional custom parser. Skip auto-detect.
+        streaming: if True (default), uses ``datasets.load_dataset(...,
+            streaming=True)`` to avoid materialising the whole dataset.
+        **load_dataset_kwargs: forwarded to ``load_dataset``.
+
+    Returns:
+        list[ConversationSample] — one per assistant tool call.
+    """
+    try:
+        from datasets import load_dataset  # type: ignore
+    except ImportError:
+        raise RuntimeError(
+            "datasets library not installed. pip install datasets, or "
+            "export the corpus to JSONL and use load_corpus_from_oai_jsonl."
+        )
+
+    model_source = model_source or f"hf:{dataset_name}"
+    ds = load_dataset(
+        dataset_name, split=split, streaming=streaming, **load_dataset_kwargs,
+    )
+
+    out: list[ConversationSample] = []
+    detected_shape: str | None = None
+    n_rows = 0
+    n_calls = 0
+    for row in ds:
+        if max_rows is not None and n_rows >= max_rows:
+            break
+        if not isinstance(row, dict):
+            continue
+        n_rows += 1
+
+        # Build the Conversation. If the caller supplied a custom
+        # parser, use it. Otherwise detect once and reuse.
+        try:
+            if row_to_conversation is not None:
+                conv = row_to_conversation(row)
+            else:
+                if detected_shape is None:
+                    detected_shape = _detect_hf_shape(row)
+                    logger.info(
+                        f"hf loader: detected shape '{detected_shape}' "
+                        f"for {dataset_name}"
+                    )
+                if detected_shape == "openai":
+                    conv = Conversation.from_dict(row)
+                elif detected_shape == "block_content":
+                    conv = _parse_block_content_row(row)
+                elif detected_shape == "conversations":
+                    conv = Conversation.from_dict({
+                        "messages": row["conversations"],
+                        "tools": row.get("tools") or [],
+                    })
+                elif detected_shape == "sharegpt":
+                    conv = _parse_sharegpt_row(row)
+                elif detected_shape == "xlam":
+                    conv = _parse_xlam_row(row)
+                else:
+                    logger.warning(
+                        f"hf loader: unknown shape for {dataset_name}; "
+                        f"pass row_to_conversation= to handle it. "
+                        f"keys={list(row.keys())[:8]}"
+                    )
+                    break
+        except Exception as e:
+            logger.debug(f"hf loader: skipping row {n_rows}: {e}")
+            continue
+        if not isinstance(conv, Conversation) or not conv.messages:
+            continue
+
+        # Detect "any assistant turn?" so reasoning-only corpora
+        # (no tool_calls) still produce training samples via the
+        # synthetic <no_tool> target when emit_no_tool_turns=True.
+        any_call_emitted = False
+        for msg_idx, msg in enumerate(conv.messages):
+            if msg.role != "assistant":
+                continue
+            if msg.tool_calls:
+                for call_idx, tc in enumerate(msg.tool_calls):
+                    if not tc.name:
+                        continue
+                    out.append(ConversationSample.from_conversation(
+                        conv=conv,
+                        prefix_len=msg_idx,
+                        tool_call=tc,
+                        tool_call_index=call_idx,
+                        model_source=model_source,
+                    ))
+                    n_calls += 1
+                    any_call_emitted = True
+            elif emit_no_tool_turns and msg.content:
+                # Plain assistant response → "<no_tool>" target. The
+                # trunk still learns from the conversation prefix; the
+                # contrastive head learns when NOT to call anything.
+                out.append(ConversationSample(
+                    conversation=list(conv.messages[:msg_idx]),
+                    tools=list(conv.tools),
+                    target_tool="<no_tool>",
+                    target_arguments="",
+                    target_tool_call_index=0,
+                    model_source=model_source,
+                ))
+                n_calls += 1
+
+    logger.info(
+        f"hf loader: {dataset_name}#{split} → {n_rows} rows / "
+        f"{n_calls} tool-call samples"
+    )
+    return out
+
+
+# ---------------------------------------------------------------------
+# Local HuggingFace cache loader — read .arrow shards directly
+# ---------------------------------------------------------------------
+
+
+def load_corpus_from_local_arrow(
+    cache_dir: Path | str,
+    model_source: str | None = None,
+    max_rows: int | None = None,
+    row_to_conversation: "Any | None" = None,
+    emit_no_tool_turns: bool = True,
+) -> list[ConversationSample]:
+    """Load ConversationSamples from a local HuggingFace cache directory.
+
+    Unlike ``load_corpus_from_hf``, this skips the ``datasets`` library
+    entirely and reads the on-disk Arrow IPC shards directly. Useful
+    when the dataset is already cached and you don't need the lib's
+    streaming/version-resolution layer.
+
+    Expected directory layout (the layout HF datasets writes by
+    default): one or more ``*.arrow`` files in ``cache_dir``, each an
+    Arrow IPC stream produced by ``Dataset.save_to_disk`` or the
+    ``json`` builder. Schema is auto-detected per file's first row
+    using the same logic as ``load_corpus_from_hf``.
+    """
+    import pyarrow as pa
+    import pyarrow.ipc as ipc
+
+    cache_dir = Path(cache_dir)
+    if not cache_dir.exists():
+        logger.warning(f"local arrow cache not found: {cache_dir}")
+        return []
+
+    shards = sorted(cache_dir.glob("*.arrow"))
+    if not shards:
+        logger.warning(f"no .arrow shards in {cache_dir}")
+        return []
+
+    model_source = model_source or f"hf-local:{cache_dir.name}"
+    out: list[ConversationSample] = []
+    detected_shape: str | None = None
+    n_rows = 0
+    n_calls = 0
+    stop = False
+
+    for shard in shards:
+        if stop:
+            break
+        try:
+            with pa.memory_map(str(shard), "r") as src:
+                rb = ipc.open_stream(src)
+                table = rb.read_all()
+        except Exception as e:
+            logger.warning(f"failed to read {shard}: {e}")
+            continue
+
+        # Convert column-major arrow to row-major dicts as we go.
+        for row in table.to_pylist():
+            if max_rows is not None and n_rows >= max_rows:
+                stop = True
+                break
+            if not isinstance(row, dict):
+                continue
+            n_rows += 1
+
+            try:
+                if row_to_conversation is not None:
+                    conv = row_to_conversation(row)
+                else:
+                    if detected_shape is None:
+                        detected_shape = _detect_hf_shape(row)
+                        logger.info(
+                            f"local arrow loader: detected shape "
+                            f"'{detected_shape}' for {cache_dir.name}"
+                        )
+                    if detected_shape == "openai":
+                        conv = Conversation.from_dict(row)
+                    elif detected_shape == "block_content":
+                        conv = _parse_block_content_row(row)
+                    elif detected_shape == "conversations":
+                        conv = Conversation.from_dict({
+                            "messages": row["conversations"],
+                            "tools": row.get("tools") or [],
+                        })
+                    elif detected_shape == "sharegpt":
+                        conv = _parse_sharegpt_row(row)
+                    elif detected_shape == "xlam":
+                        conv = _parse_xlam_row(row)
+                    elif detected_shape == "glaive":
+                        conv = _parse_glaive_row(row)
+                    else:
+                        logger.warning(
+                            f"local arrow loader: unknown shape; "
+                            f"keys={list(row.keys())[:8]}"
+                        )
+                        stop = True
+                        break
+            except Exception as e:
+                logger.debug(f"local arrow: skipping row {n_rows}: {e}")
+                continue
+
+            if not isinstance(conv, Conversation) or not conv.messages:
+                continue
+
+            # Prefer the per-row teacher model id when the dataset
+            # exposes one (Kimi-K2.5-Reasoning has meta.teacher_model;
+            # other corpora may not). Falls back to the loader-level
+            # tag. Source-conditioning works either way.
+            row_source = model_source
+            meta = row.get("meta")
+            if isinstance(meta, dict) and meta.get("teacher_model"):
+                row_source = str(meta["teacher_model"])
+
+            for msg_idx, msg in enumerate(conv.messages):
+                if msg.role != "assistant":
+                    continue
+                if msg.tool_calls:
+                    for call_idx, tc in enumerate(msg.tool_calls):
+                        if not tc.name:
+                            continue
+                        out.append(ConversationSample.from_conversation(
+                            conv=conv,
+                            prefix_len=msg_idx,
+                            tool_call=tc,
+                            tool_call_index=call_idx,
+                            model_source=row_source,
+                        ))
+                        n_calls += 1
+                elif emit_no_tool_turns and msg.content:
+                    out.append(ConversationSample(
+                        conversation=list(conv.messages[:msg_idx]),
+                        tools=list(conv.tools),
+                        target_tool="<no_tool>",
+                        target_arguments="",
+                        target_tool_call_index=0,
+                        model_source=row_source,
+                    ))
+                    n_calls += 1
+
+    logger.info(
+        f"local arrow loader: {cache_dir.name} → {n_rows} rows / "
+        f"{n_calls} samples"
+    )
+    return out
+
+
+# ---------------------------------------------------------------------
+# Glaive function-calling-v2 — custom text-embedded JSON shape
+# ---------------------------------------------------------------------
+
+_GLAIVE_ROLE_MARKER_RE = re.compile(
+    r"^(USER|ASSISTANT|FUNCTION RESPONSE):", re.MULTILINE
+)
+_GLAIVE_ROLE_NORM = {
+    "USER": "user", "ASSISTANT": "assistant", "FUNCTION RESPONSE": "tool",
+}
+
+
+def _glaive_extract_json_objects(text: str) -> list[dict[str, Any]]:
+    """Walk balanced-brace JSON objects out of free text.
+
+    Glaive's `system` field packs zero or more tool defs as separate
+    JSON objects glued together with newlines and an English preamble.
+    Standard json.loads can't parse that; we scan for top-level `{...}`
+    spans and decode each independently. Strings (including those with
+    backslash escapes) are honoured so braces inside string values
+    don't break the depth counter.
+    """
+    out: list[dict[str, Any]] = []
+    depth = 0
+    start = -1
+    in_str = False
+    i = 0
+    n = len(text)
+    while i < n:
+        c = text[i]
+        if in_str:
+            if c == "\\" and i + 1 < n:
+                i += 2
+                continue
+            if c == '"':
+                in_str = False
+        else:
+            if c == '"':
+                in_str = True
+            elif c == "{":
+                if depth == 0:
+                    start = i
+                depth += 1
+            elif c == "}":
+                depth -= 1
+                if depth == 0 and start >= 0:
+                    blob = text[start:i + 1]
+                    try:
+                        parsed = json.loads(blob)
+                        if isinstance(parsed, dict):
+                            out.append(parsed)
+                    except Exception:
+                        pass
+                    start = -1
+        i += 1
+    return out
+
+
+def _glaive_parse_functioncall(blob: str) -> dict[str, Any] | None:
+    """Extract {name, arguments} from a `<functioncall>` payload.
+
+    The payload looks like `{"name": "x", "arguments": '{...}'}` where
+    `arguments` is a single-quoted JSON string — not valid JSON. Hand-roll
+    the extraction so we don't have to rewrite the blob.
+    """
+    name_m = re.search(r'"name"\s*:\s*"([^"]+)"', blob)
+    if not name_m:
+        return None
+    name = name_m.group(1)
+    args = ""
+    # Try single-quoted arguments first (the common case)
+    args_m = re.search(
+        r'"arguments"\s*:\s*\'(.*?)\'(?=\s*[,}])', blob, re.DOTALL,
+    )
+    if args_m:
+        args = args_m.group(1)
+    else:
+        # Fallback: double-quoted string args
+        args_m2 = re.search(r'"arguments"\s*:\s*"((?:[^"\\]|\\.)*)"', blob)
+        if args_m2:
+            args = args_m2.group(1).encode().decode("unicode_escape")
+        else:
+            # Or an inline dict
+            args_m3 = re.search(
+                r'"arguments"\s*:\s*(\{.*?\})\s*[,}]', blob, re.DOTALL,
+            )
+            if args_m3:
+                args = args_m3.group(1)
+    return {"name": name, "arguments": args}
+
+
+def _parse_glaive_row(row: dict[str, Any]) -> Conversation:
+    """Parse a Glaive function-calling-v2 row → standard Conversation.
+
+    Row shape:
+      ``{"system": "<preamble + JSON tool defs>", "chat": "<USER:...
+         ASSISTANT:...<functioncall>{...}... FUNCTION RESPONSE:...>"}``
+
+    Tools come from the system field's embedded JSON; turns are split
+    on the role markers and `<functioncall>` blocks inside assistant
+    turns become ToolCalls. `<|endoftext|>` is stripped from content.
+    """
+    sys_text = str(row.get("system") or "")
+    tool_defs = _glaive_extract_json_objects(sys_text)
+    tools_norm: list[dict[str, Any]] = []
+    for t in tool_defs:
+        if not t.get("name"):
+            continue
+        tools_norm.append({"type": "function", "function": t})
+
+    chat = str(row.get("chat") or "")
+    msgs: list[dict[str, Any]] = []
+    matches = list(_GLAIVE_ROLE_MARKER_RE.finditer(chat))
+    for i, m in enumerate(matches):
+        role = _GLAIVE_ROLE_NORM.get(m.group(1))
+        if not role:
+            continue
+        start = m.end()
+        end = matches[i + 1].start() if i + 1 < len(matches) else len(chat)
+        chunk = chat[start:end].replace("<|endoftext|>", "").strip()
+
+        tool_calls: list[dict[str, Any]] = []
+        if role == "assistant":
+            # Pull each <functioncall> payload (balanced-brace) out.
+            j = 0
+            cleaned: list[str] = []
+            text_start = 0
+            while True:
+                fc = chunk.find("<functioncall>", j)
+                if fc < 0:
+                    cleaned.append(chunk[text_start:])
+                    break
+                cleaned.append(chunk[text_start:fc])
+                # Find balanced {...} after the marker
+                k = chunk.find("{", fc)
+                if k < 0:
+                    text_start = fc + len("<functioncall>")
+                    j = text_start
+                    continue
+                depth = 0
+                end_k = k
+                in_str = False
+                while end_k < len(chunk):
+                    cc = chunk[end_k]
+                    if in_str:
+                        if cc == "\\" and end_k + 1 < len(chunk):
+                            end_k += 2
+                            continue
+                        if cc == '"':
+                            in_str = False
+                    else:
+                        if cc == '"':
+                            in_str = True
+                        elif cc == "{":
+                            depth += 1
+                        elif cc == "}":
+                            depth -= 1
+                            if depth == 0:
+                                break
+                    end_k += 1
+                blob = chunk[k:end_k + 1]
+                parsed = _glaive_parse_functioncall(blob)
+                if parsed and parsed.get("name"):
+                    tool_calls.append({
+                        "type": "function",
+                        "function": {
+                            "name": str(parsed["name"]),
+                            "arguments": str(parsed.get("arguments") or ""),
+                        },
+                    })
+                text_start = end_k + 1
+                j = text_start
+
+            chunk = " ".join(p for p in (s.strip() for s in cleaned) if p)
+
+        msg_dict: dict[str, Any] = {"role": role, "content": chunk}
+        if tool_calls:
+            msg_dict["tool_calls"] = tool_calls
+        msgs.append(msg_dict)
+
+    return Conversation.from_dict({"messages": msgs, "tools": tools_norm})
+
+
+def load_corpus_from_glaive_json(
+    json_path: Path | str,
+    model_source: str | None = None,
+    max_rows: int | None = None,
+    emit_no_tool_turns: bool = True,
+) -> list[ConversationSample]:
+    """Load ConversationSamples from a Glaive function-calling JSON file.
+
+    The Glaive corpus is a single 259 MB JSON array — small enough to
+    load whole into memory. Each row is parsed via ``_parse_glaive_row``,
+    then emitted as one sample per assistant tool_call (open-vocab head
+    trains) or one sample with ``target_tool='<no_tool>'`` for
+    assistant chat turns (so the trunk still gets signal).
+    """
+    json_path = Path(json_path)
+    if not json_path.exists():
+        logger.warning(f"glaive json not found: {json_path}")
+        return []
+
+    model_source = model_source or "glaive-function-calling-v2"
+    out: list[ConversationSample] = []
+
+    try:
+        with json_path.open() as f:
+            data = json.load(f)
+    except Exception as e:
+        logger.warning(f"glaive json read failed: {e}")
+        return []
+    if not isinstance(data, list):
+        logger.warning(f"glaive json not a list: {json_path}")
+        return []
+
+    n_rows = 0
+    n_calls = 0
+    for row in data:
+        if max_rows is not None and n_rows >= max_rows:
+            break
+        if not isinstance(row, dict):
+            continue
+        n_rows += 1
+        try:
+            conv = _parse_glaive_row(row)
+        except Exception as e:
+            logger.debug(f"glaive: skipping row {n_rows}: {e}")
+            continue
+        if not isinstance(conv, Conversation) or not conv.messages:
+            continue
+
+        for msg_idx, msg in enumerate(conv.messages):
+            if msg.role != "assistant":
+                continue
+            if msg.tool_calls:
+                for call_idx, tc in enumerate(msg.tool_calls):
+                    if not tc.name:
+                        continue
+                    out.append(ConversationSample.from_conversation(
+                        conv=conv,
+                        prefix_len=msg_idx,
+                        tool_call=tc,
+                        tool_call_index=call_idx,
+                        model_source=model_source,
+                    ))
+                    n_calls += 1
+            elif emit_no_tool_turns and msg.content:
+                out.append(ConversationSample(
+                    conversation=list(conv.messages[:msg_idx]),
+                    tools=list(conv.tools),
+                    target_tool="<no_tool>",
+                    target_arguments="",
+                    target_tool_call_index=0,
+                    model_source=model_source,
+                ))
+                n_calls += 1
+
+    logger.info(
+        f"glaive json loader: {json_path.name} → {n_rows} rows / "
+        f"{n_calls} samples"
+    )
+    return out
+
+
+# ---------------------------------------------------------------------
+# Forge feature-outcome corpus  (ADR 0006 Phase A)
+# ---------------------------------------------------------------------
+
+
+def load_corpus_from_feature_outcomes(
+    path: Path = Path("data/nn/feature_outcomes.jsonl"),
+    dedupe_by_feature: bool = True,
+) -> list[ConversationSample]:
+    """Convert each Forge feature-outcome row into one ConversationSample.
+
+    Schema (written by the orchestrator):
+        feature_text, model_used, success (bool), n_turns, n_tool_calls,
+        duration_s, project_id, feature_id, session_id, extras{}.
+
+    The feature text becomes the only "message" (one-turn batch); the
+    model id flows through `model_source` so the trunk's source-
+    embedding can bias predictions per teacher model. All other state
+    (mood, tool history, images) is zero — Forge predictions are made
+    *before* a session runs, so we deliberately do not condition on
+    in-turn signals that don't exist yet at prediction time.
+
+    Other heads' targets stay None — those heads are skipped for these
+    samples (HeadSpec.optional_target). The only label this corpus
+    carries is `target_feature_success`.
+
+    ``dedupe_by_feature`` (default True) collapses retries — every
+    (project_id, feature_id) pair contributes ONE sample with the
+    final outcome (last row by ts). Without this, a stuck feature that
+    retried 200 times before being parked dominates the corpus and
+    pushes the head to "always predict fail". The orchestrator's
+    decision question is "is this feature solvable?", not "will this
+    particular retry succeed?", so the final outcome is the right
+    label. Set False if you actually want per-session granularity
+    (e.g. for predict_failure_repeats analysis).
+
+    Rows with missing required fields are dropped silently. Skipping
+    bad rows is the right behaviour because the JSONL is append-only
+    and a single malformed write must not bork the next training run.
+    """
+    if not path.exists():
+        return []
+    try:
+        lines = path.read_text(encoding="utf-8").splitlines()
+    except OSError:
+        return []
+
+    raw_rows: list[dict[str, Any]] = []
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            row = json.loads(line)
+        except Exception:
+            continue
+        text = (row.get("feature_text") or "").strip()
+        success = row.get("success")
+        if not text or success is None:
+            continue
+        raw_rows.append(row)
+
+    if dedupe_by_feature:
+        # Keep the latest row per (project_id, feature_id). If feature_id
+        # is missing (older rows) we fall back to (project_id, feature_text)
+        # so a feature with no id still gets one slot, not one slot per
+        # retry.
+        latest: dict[tuple, dict[str, Any]] = {}
+        for row in raw_rows:
+            key = (
+                row.get("project_id"),
+                row.get("feature_id") or row.get("feature_text"),
+            )
+            prev = latest.get(key)
+            if prev is None or row.get("ts", 0) > prev.get("ts", 0):
+                latest[key] = row
+        raw_rows = list(latest.values())
+
+    out: list[ConversationSample] = []
+    for row in raw_rows:
+        text = (row.get("feature_text") or "").strip()
+        success = bool(row.get("success"))
+        # Feature description as a one-turn user message — the model's
+        # job here is "given just the feature text, predict success".
+        # No tool calls, no tool list — feature_success is the only
+        # head this row labels.
+        out.append(ConversationSample(
+            conversation=[ChatMessage(role="user", content=text)],
+            tools=[],
+            target_tool="<no_tool>",
+            target_value=1.0 if success else 0.0,
+            model_source=row.get("model_used") or "<unknown>",
+            target_feature_success=1.0 if success else 0.0,
+        ))
+    return out
+
+
+# ---------------------------------------------------------------------
+# Replay buffer
+# ---------------------------------------------------------------------
+
+
+@dataclass
+class ReplayBuffer:
+    """Circular buffer for online training samples."""
+
+    capacity: int = 8192
+    _buf: collections.deque = field(default_factory=lambda: collections.deque(maxlen=8192))
+
+    def __post_init__(self) -> None:
+        # Re-bind maxlen to honor the constructor arg
+        if self._buf.maxlen != self.capacity:
+            self._buf = collections.deque(self._buf, maxlen=self.capacity)
+
+    def push(self, sample: ConversationSample) -> None:
+        self._buf.append(sample)
+
+    def __len__(self) -> int:
+        return len(self._buf)
+
+    def sample(self, n: int) -> list[ConversationSample]:
+        if n <= 0 or not self._buf:
+            return []
+        buf = list(self._buf)
+        # Surprise-weighted sampling (Prioritized Experience Replay):
+        # samples Caudate guessed badly on appear more often. Weight is
+        # clamped above zero so easy samples still occasionally show up
+        # (otherwise the model forgets what it already knows). Sampling
+        # is with replacement, so n may exceed buffer size — high-
+        # priority items can repeat within a batch.
+        weights = [max(0.05, float(getattr(s, "surprise", 0.5))) for s in buf]
+        return random.choices(buf, weights=weights, k=n)
+
+    def all(self) -> list[ConversationSample]:
+        return list(self._buf)
+
+
+# ---------------------------------------------------------------------
+# Batch collation
+# ---------------------------------------------------------------------
+
+
+_NO_TOOL_DESCRIPTION = (
+    "answer the user directly without calling any tool"
+)
+
+
+def _format_tool_spec(t: "ToolDef") -> str:
+    """Render a ToolDef as the "{name}: {description}" string that the
+    contrastive tool head's text encoder consumes. Capped at 240 chars
+    to keep batched embedding tractable on big tool registries."""
+    name = (t.name or "").strip()
+    desc = (t.description or "").strip()
+    if desc:
+        return f"{name}: {desc}"[:240]
+    return name[:240]
+
+
+def collate(
+    samples: list[ConversationSample],
+    cfg: NNConfig,
+    vocab: ToolVocab,
+    source_vocab: "SourceVocab | None" = None,
+) -> dict[str, Any]:
+    """Turn a list of ConversationSamples into model-ready tensors.
+
+    The chat-schema conversation gets flattened into role-prefixed
+    strings by ``conversation_to_strings``; tool history is derived
+    from past assistant tool_calls. Tool candidates per sample are
+    surfaced as ``tool_specs`` (a list of "{name}: {description}"
+    strings) so the contrastive tool head can score them. The first
+    candidate slot is always a synthetic ``<no_tool>`` entry — that's
+    how "the assistant answered directly without calling anything"
+    becomes a real class the head can predict.
+    """
+    B = len(samples)
+    msg_window = cfg.msg_window
+    hist_window = cfg.history_window
+    img_window = cfg.image_window if cfg.use_vision else 0
+    K = cfg.max_tools_per_sample
+
+    messages: list[list[str]] = []
+    image_paths: list[list[str]] = []
+    tool_ids = torch.zeros((B, hist_window), dtype=torch.long)
+    mood = torch.zeros((B, cfg.mood_dim), dtype=torch.float32)
+
+    target_tier = torch.zeros(B, dtype=torch.long)
+    target_think = torch.zeros(B, dtype=torch.float32)
+    target_value = torch.zeros(B, dtype=torch.float32)
+    source_id = torch.zeros(B, dtype=torch.long)
+
+    # Contrastive tool head inputs.
+    # tool_specs: per-sample list of "{name}: {description}" strings,
+    #   first slot is always "<no_tool>: answer directly..."
+    # target_tool_idx: index of the called tool within that list, or -1
+    #   if the called tool isn't among the candidates (skip-this-sample).
+    tool_specs: list[list[str]] = []
+    target_tool_idx = torch.full((B,), -1, dtype=torch.long)
+    tool_target_valid = torch.zeros(B, dtype=torch.bool)
+
+    for i, s in enumerate(samples):
+        # Flatten the role-tagged conversation into text the encoder
+        # already speaks. Pad/truncate to msg_window.
+        msg_strs = conversation_to_strings(s.conversation)[-msg_window:]
+        if len(msg_strs) < msg_window:
+            msg_strs = [""] * (msg_window - len(msg_strs)) + msg_strs
+        messages.append(msg_strs)
+
+        if cfg.use_vision:
+            imgs = list(s.image_paths)[-img_window:]
+            if len(imgs) < img_window:
+                imgs = [""] * (img_window - len(imgs)) + imgs
+            image_paths.append(imgs)
+        else:
+            image_paths.append([])
+
+        # Tool history = the sequence of tool names called in this
+        # conversation prefix. Stable ordering preserves recency.
+        # Clamp every id below tool_vocab_size — the encoder's tool
+        # embedding table is fixed-size and a vocab that grows past
+        # it (public corpora can hit 10K+ tools) would otherwise
+        # crash with index-out-of-range. Modulo wraps high ids into
+        # the table with some collisions; better than missing data.
+        tool_hist = conversation_tool_history(s.conversation)[-hist_window:]
+        hist = [vocab.get(t) % cfg.tool_vocab_size for t in tool_hist]
+        if len(hist) < hist_window:
+            hist = [vocab.SPECIAL["<pad>"]] * (hist_window - len(hist)) + hist
+        tool_ids[i] = torch.tensor(hist, dtype=torch.long)
+
+        mood[i] = torch.tensor(s.mood[: cfg.mood_dim], dtype=torch.float32)
+        target_tier[i] = int(s.target_tier)
+        target_think[i] = float(s.target_think)
+        target_value[i] = float(s.target_value)
+        if source_vocab is not None:
+            # `add` is idempotent — first occurrence assigns an id, later
+            # ones look it up. Falls back to <unknown> if the cap fills.
+            source_id[i] = source_vocab.add(getattr(s, "model_source", "<unknown>"))
+        else:
+            source_id[i] = 0  # <unknown> — keeps the model in baseline mode
+
+        # Build the candidate tool list for this sample. The synthetic
+        # <no_tool> entry occupies slot 0 so the head always has a
+        # "don't call anything" option available — without it the
+        # softmax would be forced to nominate something even when the
+        # right answer is "respond directly".
+        sample_tools = list(s.tools)[: max(0, K - 1)]
+        specs = [f"<no_tool>: {_NO_TOOL_DESCRIPTION}"]
+        for t in sample_tools:
+            specs.append(_format_tool_spec(t))
+        tool_specs.append(specs)
+
+        # Resolve target_tool to an index in `specs`. Match by name.
+        target_name = (s.target_tool or "<no_tool>")
+        if target_name == "<no_tool>" or target_name == "":
+            target_tool_idx[i] = 0
+            tool_target_valid[i] = True
+        else:
+            found = -1
+            for k, t in enumerate(sample_tools):
+                if t.name == target_name:
+                    found = k + 1  # +1 for the <no_tool> prefix slot
+                    break
+            if found >= 0:
+                target_tool_idx[i] = found
+                tool_target_valid[i] = True
+            # Else: target_tool isn't among candidates → leave
+            # target_tool_idx=-1, tool_target_valid=False; the trainer
+            # skips this sample's contribution to the tool loss.
+
+    out = {
+        "messages": messages,
+        "image_paths": image_paths,
+        "tool_ids": tool_ids,
+        "mood": mood,
+        "target_tier": target_tier,
+        "target_think": target_think,
+        "target_value": target_value,
+        "source_id": source_id,
+        # Contrastive tool head inputs/targets
+        "tool_specs": tool_specs,
+        "target_tool_idx": target_tool_idx,
+        "tool_target_valid": tool_target_valid,
+    }
+
+    # Extended-head targets — only emit if EVERY sample in this batch
+    # has the label. A partially-labeled batch would have to use a fake
+    # value (e.g. 0.0) for the unlabeled rows, which would train the
+    # head against a synthetic "no" signal and corrupt learning. So
+    # if even one sample is missing the label, we drop the field
+    # entirely and the trainer skips that head for this batch (via
+    # HeadSpec.optional_target).
+    _SCALAR_BCE_FIELDS = (
+        "target_memory_write", "target_cache_hit", "target_permission",
+        "target_refusal", "target_code_response", "target_stall",
+        "target_stop_iter", "target_compaction",
+        "target_subagent_spawn",
+        "target_feature_success",
+    )
+    _SCALAR_MSE_FIELDS = (
+        "target_latency_s", "target_token_budget", "target_reward_model",
+    )
+    for field_name in _SCALAR_BCE_FIELDS + _SCALAR_MSE_FIELDS:
+        vals = [getattr(s, field_name, None) for s in samples]
+        if all(v is not None for v in vals):
+            out[field_name] = torch.tensor(
+                [float(v) for v in vals], dtype=torch.float32
+            )
+
+    # difficulty is a 3-class CE target → int64
+    diffs = [getattr(s, "target_difficulty", None) for s in samples]
+    if all(d is not None for d in diffs):
+        out["target_difficulty"] = torch.tensor(
+            [int(d) for d in diffs], dtype=torch.long
+        )
+
+    # mood_pred is a 4-D vector
+    moods = [getattr(s, "target_mood_pred", None) for s in samples]
+    if all(m is not None for m in moods):
+        out["target_mood_pred"] = torch.tensor(
+            [list(m) + [0.0]*(4-len(m)) if len(m) < 4 else list(m)[:4] for m in moods],
+            dtype=torch.float32,
+        )
+    return out
+
+
+def split_train_eval(
+    samples: list[ConversationSample], eval_ratio: float, seed: int,
+) -> tuple[list[ConversationSample], list[ConversationSample]]:
+    rng = random.Random(seed)
+    shuffled = list(samples)
+    rng.shuffle(shuffled)
+    n_eval = max(1, int(len(shuffled) * eval_ratio)) if shuffled else 0
+    return shuffled[n_eval:], shuffled[:n_eval]
+
+
+def iter_batches(
+    samples: list[ConversationSample],
+    batch_size: int,
+    cfg: NNConfig,
+    vocab: ToolVocab,
+    shuffle: bool = True,
+    source_vocab: "SourceVocab | None" = None,
+) -> Iterator[dict[str, Any]]:
+    if shuffle:
+        samples = list(samples)
+        random.shuffle(samples)
+    for i in range(0, len(samples), batch_size):
+        chunk = samples[i:i + batch_size]
+        if not chunk:
+            continue
+        yield collate(chunk, cfg, vocab, source_vocab=source_vocab)