caudate-cli 0.1.0__py3-none-any.whl

nn/format.py

@@ -0,0 +1,235 @@
+"""Standard chat-conversation schema for Caudate.
+
+The on-disk and in-memory training format mirrors the OpenAI tool-call
+JSON shape that every public tool-use dataset uses (xLAM, ToolBench,
+ToolAlpaca, Gorilla, API-Bank). This lets Caudate train directly on
+external corpora without an adapter step.
+
+Schema (one training example):
+
+    {
+      "conversations": [
+        {"role": "system" | "user" | "assistant" | "tool",
+         "content": str,
+         "tool_calls"?: [
+            {"type": "function",
+             "function": {"name": str, "arguments": str (JSON)}}
+         ],
+         "name"?: str,            # required when role == "tool"
+         "tool_call_id"?: str},   # links a tool result back to the call
+        ...
+      ],
+      "tools": [
+        {"type": "function",
+         "function": {"name": str, "description": str,
+                      "parameters": dict (JSON schema)}}
+      ]
+    }
+
+These dataclasses are deliberately thin — no validation beyond what
+parsing needs. The trainer treats malformed rows as skip-and-continue
+rather than raising, because external datasets are noisy.
+
+Cognos-specific signals (mood, the 14 optional heads, model_source,
+images) live on ConversationSample in nn/data.py, NOT in this schema —
+that's the boundary that keeps external data ingestible.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+
+# ─── Tool call (assistant → tool) ─────────────────────────────────────
+
+
+@dataclass
+class ToolCall:
+    """One function call inside an assistant message."""
+
+    name: str
+    arguments: str = ""        # JSON-encoded string per OpenAI spec
+    id: str | None = None      # tool_call_id; optional in some datasets
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "ToolCall":
+        # OpenAI shape: {"type": "function", "function": {"name", "arguments"}}
+        # Some datasets flatten: {"name", "arguments"}. Accept both.
+        fn = d.get("function") if isinstance(d.get("function"), dict) else d
+        args = fn.get("arguments", "")
+        # Arguments can be a dict in some datasets — normalise to JSON string
+        # so downstream consumers don't have to handle both shapes.
+        if isinstance(args, dict):
+            args = json.dumps(args)
+        return cls(
+            name=str(fn.get("name") or ""),
+            arguments=str(args or ""),
+            id=d.get("id") or d.get("tool_call_id"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {
+            "type": "function",
+            "function": {"name": self.name, "arguments": self.arguments},
+        }
+        if self.id is not None:
+            out["id"] = self.id
+        return out
+
+
+# ─── Chat message ─────────────────────────────────────────────────────
+
+
+@dataclass
+class ChatMessage:
+    """One conversation turn. Roles: system | user | assistant | tool."""
+
+    role: str
+    content: str = ""
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    name: str | None = None          # tool name when role == "tool"
+    tool_call_id: str | None = None  # links tool result to the call
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "ChatMessage":
+        # `content` is sometimes a list of content-blocks (OpenAI vision
+        # shape). Flatten to text — Caudate's encoder doesn't currently
+        # consume content-block structure, only the text. Vision still
+        # works through image_paths on ConversationSample.
+        content = d.get("content")
+        if isinstance(content, list):
+            chunks: list[str] = []
+            for blk in content:
+                if isinstance(blk, dict):
+                    txt = blk.get("text") or blk.get("content") or ""
+                    if txt:
+                        chunks.append(str(txt))
+                elif isinstance(blk, str):
+                    chunks.append(blk)
+            content = " ".join(chunks)
+        elif content is None:
+            content = ""
+        raw_calls = d.get("tool_calls") or []
+        calls = [ToolCall.from_dict(c) for c in raw_calls if isinstance(c, dict)]
+        return cls(
+            role=str(d.get("role") or ""),
+            content=str(content),
+            tool_calls=calls,
+            name=d.get("name"),
+            tool_call_id=d.get("tool_call_id"),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {"role": self.role, "content": self.content}
+        if self.tool_calls:
+            out["tool_calls"] = [tc.to_dict() for tc in self.tool_calls]
+        if self.name is not None:
+            out["name"] = self.name
+        if self.tool_call_id is not None:
+            out["tool_call_id"] = self.tool_call_id
+        return out
+
+
+# ─── Tool definition (what's available for the assistant to call) ─────
+
+
+@dataclass
+class ToolDef:
+    """One tool entry in the `tools` field of a training example."""
+
+    name: str
+    description: str = ""
+    parameters: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "ToolDef":
+        fn = d.get("function") if isinstance(d.get("function"), dict) else d
+        params = fn.get("parameters") or {}
+        if not isinstance(params, dict):
+            params = {}
+        return cls(
+            name=str(fn.get("name") or ""),
+            description=str(fn.get("description") or ""),
+            parameters=params,
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": self.name,
+                "description": self.description,
+                "parameters": self.parameters,
+            },
+        }
+
+
+# ─── Conversation (one training example, schema-side only) ───────────
+
+
+@dataclass
+class Conversation:
+    """A single training example in standard chat-tool-call shape.
+
+    This is the disk/serialization view. The training-time view —
+    Conversation plus Cognos-specific signals (mood, model_source,
+    image_paths, optional head targets) — is ConversationSample in
+    nn/data.py."""
+
+    messages: list[ChatMessage] = field(default_factory=list)
+    tools: list[ToolDef] = field(default_factory=list)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> "Conversation":
+        # Accept both "messages" (modern) and "conversations" (some legacy
+        # datasets) as the field name for the turn list.
+        raw_msgs = d.get("messages")
+        if raw_msgs is None:
+            raw_msgs = d.get("conversations") or []
+        raw_tools = d.get("tools") or []
+        return cls(
+            messages=[
+                ChatMessage.from_dict(m) for m in raw_msgs if isinstance(m, dict)
+            ],
+            tools=[
+                ToolDef.from_dict(t) for t in raw_tools if isinstance(t, dict)
+            ],
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "messages": [m.to_dict() for m in self.messages],
+            "tools": [t.to_dict() for t in self.tools],
+        }
+
+
+# ─── Iteration helpers used by ConversationSample & loaders ──────────
+
+
+def iter_assistant_tool_calls(
+    conv: Conversation,
+) -> list[tuple[int, ToolCall]]:
+    """Yield (message_index, ToolCall) for every assistant tool call.
+
+    Used by loaders that need to emit one training sample *per tool
+    call* — the canonical Caudate training granularity is "given this
+    prefix, what tool gets called next?". Each tool call inside an
+    assistant message is one training row."""
+    out: list[tuple[int, ToolCall]] = []
+    for i, msg in enumerate(conv.messages):
+        if msg.role != "assistant":
+            continue
+        for tc in msg.tool_calls:
+            out.append((i, tc))
+    return out
+
+
+__all__ = [
+    "ChatMessage",
+    "ToolCall",
+    "ToolDef",
+    "Conversation",
+    "iter_assistant_tool_calls",
+]

nn/heads.py

@@ -0,0 +1,432 @@
+"""Caudate head registry — pluggable output specialists on a shared trunk.
+
+Caudate is a single neural network. The transformer trunk produces one
+pooled CLS representation per turn; multiple *heads* on top of that
+representation each predict one categorical or scalar target. Each head
+is one **routing job** (which tool, which tier, should-think, expected
+reward, and — added in this refactor — preference, memory-write, cache,
+permission, …).
+
+The registry pattern means:
+
+  - Adding a new routing job = appending one `HeadSpec` to a list. No
+    surgery in `Caudate` itself.
+  - The trainer iterates the spec list to compute one loss per head and
+    aggregate them by `loss_weight`.
+  - State-dict layout matches the legacy `_Heads` module exactly
+    (`heads.tool.*`, `heads.tier.*`, …) so existing `caudate.pt` files
+    load cleanly with `strict=False` — new heads get fresh init,
+    existing heads keep their trained weights.
+
+Adding a head is *not* a forking-into-a-second-model operation. The
+trunk weights keep absorbing gradient from every head, so each new
+job makes the shared representation richer for the existing jobs too.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Literal
+
+import torch
+import torch.nn as nn
+
+LossType = Literal["ce", "bce", "mse_sigmoid"]
+
+
+@dataclass(frozen=True)
+class HeadSpec:
+    """One routing job on the shared trunk.
+
+    `name` is the module name in `state_dict` (`heads.<name>.0.weight`,
+    so picking it well matters for backward compat). `output_key` is
+    the key under which forward() returns this head's tensor — kept
+    distinct from `name` only because the legacy heads used different
+    conventions (`tool_logits` vs the new `value`/`think_logit`).
+    """
+
+    name: str                       # module name (state_dict key)
+    out_dim: int                    # output dimension
+    output_key: str                 # key in forward()'s output dict
+    loss_type: LossType             # ce | bce | mse_sigmoid
+    loss_weight: float = 1.0        # weight in total loss
+    hidden_ratio: float = 0.5       # hidden = round(d_model * hidden_ratio)
+    target_field: str = ""          # batch field for the target; defaults to f"target_{name}"
+    optional_target: bool = False   # if True, missing target in batch → skip this head's loss
+
+    @property
+    def target(self) -> str:
+        return self.target_field or f"target_{self.name}"
+
+
+# ---------------------------------------------------------------------
+# Standard head bundle: the 4 original Caudate routing jobs.
+# ---------------------------------------------------------------------
+
+STANDARD_HEADS: tuple[HeadSpec, ...] = (
+    # NOTE: the `tool` head used to live here as a fixed-vocab CE
+    # classifier (out_dim=32). Phase-2 of the format migration replaced
+    # it with `ContrastiveToolHead` — see below — which scores candidate
+    # tools by description embedding rather than fixed slot. The tool
+    # head is NOT in the standard HeadSpec registry anymore because it
+    # needs custom forward + loss paths (variable per-sample candidate
+    # count, masked contrastive softmax). The trainer wires it
+    # separately.
+    HeadSpec(name="tier",  out_dim=2,  output_key="tier_logits",
+             loss_type="ce",  loss_weight=0.5, hidden_ratio=0.5),
+    HeadSpec(name="think", out_dim=1,  output_key="think_logit",
+             loss_type="bce", loss_weight=0.2, hidden_ratio=0.25),
+    HeadSpec(name="value", out_dim=1,  output_key="value",
+             loss_type="mse_sigmoid", loss_weight=0.5, hidden_ratio=0.25),
+)
+
+# ---------------------------------------------------------------------
+# Extended heads — new routing jobs on the same brain.
+#
+# These are wired structurally now (D-foundation); their label
+# collectors come online in D-heads. Each carries `optional_target=True`
+# so a batch missing the target field just skips that head's loss
+# instead of crashing — important for incremental rollout.
+# ---------------------------------------------------------------------
+
+EXTENDED_HEADS: tuple[HeadSpec, ...] = (
+    # --- Original D-heads (memory / cache / permission) -----------------
+    HeadSpec(name="memory_write", out_dim=1, output_key="memory_write_logit",
+             loss_type="bce", loss_weight=0.15, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="cache_hit", out_dim=1, output_key="cache_hit_logit",
+             loss_type="bce", loss_weight=0.15, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="permission", out_dim=1, output_key="permission_logit",
+             loss_type="bce", loss_weight=0.15, hidden_ratio=0.25,
+             optional_target=True),
+
+    # --- Tier 1: response-shape predictors ------------------------------
+    # Each head pre-classifies the *shape* of the upcoming response so
+    # routing can act on the prediction instead of waiting for it.
+    HeadSpec(name="refusal", out_dim=1, output_key="refusal_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="code_response", out_dim=1, output_key="code_response_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="stall", out_dim=1, output_key="stall_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="difficulty", out_dim=3, output_key="difficulty_logits",
+             loss_type="ce", loss_weight=0.15, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="stop_iter", out_dim=1, output_key="stop_iter_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="compaction", out_dim=1, output_key="compaction_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+
+    # --- Tier 2: continuous + multi-output predictors -------------------
+    HeadSpec(name="latency_s", out_dim=1, output_key="latency_s",
+             loss_type="mse_sigmoid", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="token_budget", out_dim=1, output_key="token_budget",
+             loss_type="mse_sigmoid", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    # 4-D mood (frustration, curiosity, satisfaction, urgency) — sigmoid
+    # so each axis is in [0,1] and they're independent (not softmax).
+    HeadSpec(name="mood_pred", out_dim=4, output_key="mood_pred",
+             loss_type="mse_sigmoid", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+    HeadSpec(name="subagent_spawn", out_dim=1, output_key="subagent_spawn_logit",
+             loss_type="bce", loss_weight=0.10, hidden_ratio=0.25,
+             optional_target=True),
+
+    # --- Forge feature-success head (ADR 0006 Phase A) ------------------
+    # Predicts P(feature succeeds | feature_text, model_source) for the
+    # Forge orchestrator. Label source is `data/nn/feature_outcomes.jsonl`
+    # — one row per completed (success or failure) Forge feature. The
+    # trunk's source-embedding already captures per-model bias, so this
+    # head only needs to read the pooled CLS conditioned on the feature
+    # text. hidden_ratio=0.5 because the signal is concentrated in the
+    # text+source pair, not in the broader turn-state.
+    HeadSpec(name="feature_success", out_dim=1, output_key="feature_success_logit",
+             loss_type="bce", loss_weight=0.20, hidden_ratio=0.5,
+             optional_target=True),
+
+    # --- Tier 4b: reward-model head -------------------------------------
+    # Scores the quality of an external draft (replaces the heuristic
+    # `_score_draft` in api/openai_compat.py once trained). Same shape
+    # as `value` but trained on arbitration pair labels. Sees pooled CLS
+    # of the turn-state, which is enough for "given this prompt, how
+    # likely is a strong response?" — the more precise per-draft scoring
+    # waits for the generative-decoder phase.
+    HeadSpec(name="reward_model", out_dim=1, output_key="reward_model_logit",
+             loss_type="mse_sigmoid", loss_weight=0.20, hidden_ratio=0.5,
+             optional_target=True),
+
+    # --- Tier 3: self-supervised auxiliaries ----------------------------
+    # next_prompt — predicts an embedding of the *next* user message
+    # given the current pooled CLS. Cosine-similarity loss against the
+    # encoder's own embedding of the actual next message (when it
+    # arrives). Auxiliary task that lifts trunk quality without any
+    # external label. Output is in d_model-dim CLS space so the loss
+    # is a direct cosine alignment.
+    # Note: out_dim matches the trunk's d_model (256). The mse_sigmoid
+    # loss isn't quite right for cosine alignment, so the trainer has a
+    # bespoke loss path for this head — see _aux_next_prompt_loss.
+    HeadSpec(name="next_prompt", out_dim=256, output_key="next_prompt_embed",
+             loss_type="mse_sigmoid", loss_weight=0.05, hidden_ratio=1.0,
+             optional_target=True),
+)
+
+ALL_HEADS: tuple[HeadSpec, ...] = STANDARD_HEADS + EXTENDED_HEADS
+
+
+# ---------------------------------------------------------------------
+# Module
+# ---------------------------------------------------------------------
+
+
+def _make_head(d_model: int, spec: HeadSpec, dropout: float) -> nn.Module:
+    hidden = max(1, int(d_model * spec.hidden_ratio))
+    if spec.hidden_ratio >= 1.0:
+        # Wider heads (tool) include dropout to control overfitting on
+        # the larger fan-in.
+        return nn.Sequential(
+            nn.Linear(d_model, hidden),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden, spec.out_dim),
+        )
+    return nn.Sequential(
+        nn.Linear(d_model, hidden),
+        nn.GELU(),
+        nn.Linear(hidden, spec.out_dim),
+    )
+
+
+class HeadRegistry(nn.ModuleDict):
+    """ModuleDict of named heads + the spec list driving the forward pass.
+
+    Inheriting from `ModuleDict` (rather than holding one as an attr)
+    keeps the state-dict layout flat (`heads.<name>.*`) so legacy
+    checkpoints load cleanly.
+    """
+
+    def __init__(self, d_model: int, dropout: float, specs: tuple[HeadSpec, ...]):
+        super().__init__({s.name: _make_head(d_model, s, dropout) for s in specs})
+        # `_specs` is a plain Python attr (not a Module) — won't pollute state_dict
+        object.__setattr__(self, "_specs", tuple(specs))
+
+        # Xavier init on linear layers; zero bias.
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    @property
+    def specs(self) -> tuple[HeadSpec, ...]:
+        return self._specs   # type: ignore[attr-defined]
+
+    def forward(self, cls: torch.Tensor) -> dict[str, torch.Tensor]:
+        out: dict[str, torch.Tensor] = {}
+        for spec in self.specs:
+            y = self[spec.name](cls)
+            if spec.out_dim == 1:
+                y = y.squeeze(-1)   # (B, 1) → (B,)
+            out[spec.output_key] = y
+        return out
+
+
+# ---------------------------------------------------------------------
+# Loss helpers — used by the trainer
+# ---------------------------------------------------------------------
+
+
+def head_loss(spec: HeadSpec, output: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    if spec.loss_type == "ce":
+        return nn.functional.cross_entropy(output, target.long())
+    if spec.loss_type == "bce":
+        return nn.functional.binary_cross_entropy_with_logits(output, target.float())
+    if spec.loss_type == "mse_sigmoid":
+        return nn.functional.mse_loss(torch.sigmoid(output), target.float())
+    raise ValueError(f"unknown loss_type {spec.loss_type!r} for head {spec.name!r}")
+
+
+def head_accuracy(spec: HeadSpec, output: torch.Tensor, target: torch.Tensor) -> float | None:
+    """Return classification accuracy for ce/bce; None for regression heads."""
+    with torch.no_grad():
+        if spec.loss_type == "ce":
+            return (output.argmax(-1) == target.long()).float().mean().item()
+        if spec.loss_type == "bce":
+            return ((output > 0).float() == target.float()).float().mean().item()
+        return None
+
+
+# ---------------------------------------------------------------------
+# Contrastive tool head — open-vocab, replaces the fixed 32-slot
+# classifier. Phase 2 of the format migration.
+# ---------------------------------------------------------------------
+
+
+class ContrastiveToolHead(nn.Module):
+    """Open-vocabulary tool predictor.
+
+    Architecture: pooled CLS → ``query`` vector. Each candidate tool's
+    "{name}: {description}" string → ``tool_emb`` via the shared text
+    encoder, projected into the same space. Score = query · tool_emb;
+    softmax over the per-sample candidate set yields the tool
+    distribution. Invalid slots (when a sample has fewer than
+    ``max_tools_per_sample`` candidates) are masked to -inf so they
+    contribute zero probability.
+
+    Why contrastive instead of fixed-slot:
+      - Open vocabulary — train on 80K xLAM tools, generalize to tools
+        never seen at training time by reading their description.
+      - Variable candidate sets per turn — different sessions expose
+        different tool registries; the head respects that.
+      - Zero "tool collision" problem — the old vocab-mod-N approach
+        would conflate distinct tools onto the same slot once vocab
+        size > n_tools_out.
+
+    Why this isn't a HeadSpec:
+      - Its forward needs the per-sample candidate list (not just the
+        pooled CLS), so the generic HeadRegistry signature doesn't fit.
+      - Its loss is a contrastive softmax with masking, not the
+        standard CE/BCE/MSE path the generic loss helper supports.
+
+    The trainer wires this head explicitly. ``Caudate.forward``
+    returns its logits under ``"tool_logits"`` and its mask under
+    ``"tool_mask"`` so downstream paths can mix accuracy and
+    confidence inspection.
+    """
+
+    def __init__(self, cfg, text_embedder):
+        super().__init__()
+        self.cfg = cfg
+        # Query projection: pooled CLS → tool-embedding space
+        d_out = cfg.n_tools_out          # now: tool embedding dim, not # of slots
+        self.query_proj = nn.Sequential(
+            nn.Linear(cfg.d_model, cfg.d_model),
+            nn.GELU(),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.d_model, d_out),
+        )
+        # Tool description projection: text encoder output → same space
+        self.tool_proj = nn.Sequential(
+            nn.Linear(cfg.text_embed_dim, cfg.d_model),
+            nn.GELU(),
+            nn.Linear(cfg.d_model, d_out),
+        )
+        # Shared with the state encoder; we don't own it, we use it.
+        # Frozen sentence-transformer (or hash-embed fallback offline).
+        self._text_embedder = text_embedder
+        # Learned temperature so the softmax sharpness can adapt.
+        self.log_temperature = nn.Parameter(torch.zeros(1))
+
+        # Init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+    @torch.no_grad()
+    def _encode_descriptions(self, flat_texts: list[str]) -> torch.Tensor:
+        """Encode a flat list of tool description strings into the
+        shared text-embedding space. Frozen — no gradient flows back
+        through the sentence-transformer.
+        """
+        if not flat_texts:
+            return torch.zeros(
+                (0, self.cfg.text_embed_dim),
+                device=next(self.parameters()).device,
+            )
+        return self._text_embedder.encode(flat_texts)
+
+    def embed_tools(
+        self,
+        tool_specs: list[list[str]],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode per-sample candidate tool strings into projected
+        embeddings + a validity mask.
+
+        Args:
+            tool_specs: B-sized list, each entry is a list of
+                "{name}: {description}" strings (zero-or-more).
+
+        Returns:
+            tool_emb: (B, K, n_tools_out) zero-padded
+            mask:     (B, K) bool — True at valid slots
+        """
+        device = next(self.parameters()).device
+        B = len(tool_specs)
+        K = self.cfg.max_tools_per_sample
+
+        # Flatten across samples for one batched embedder call.
+        flat: list[str] = []
+        per_sample_count: list[int] = []
+        for ts in tool_specs:
+            kept = list(ts)[:K]
+            per_sample_count.append(len(kept))
+            flat.extend(kept)
+
+        if flat:
+            text_embs = self._encode_descriptions(flat).to(device)
+            proj = self.tool_proj(text_embs)        # (N_flat, n_tools_out)
+        else:
+            proj = torch.zeros((0, self.cfg.n_tools_out), device=device)
+
+        out = torch.zeros(B, K, self.cfg.n_tools_out, device=device)
+        mask = torch.zeros(B, K, dtype=torch.bool, device=device)
+        cursor = 0
+        for i, k in enumerate(per_sample_count):
+            if k > 0:
+                out[i, :k] = proj[cursor:cursor + k]
+                mask[i, :k] = True
+            cursor += k
+        return out, mask
+
+    def forward(
+        self,
+        pooled_cls: torch.Tensor,
+        tool_emb: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute (B, K) logits over candidate tools. Invalid slots
+        are returned as a very negative number so any softmax over
+        the result effectively ignores them."""
+        query = self.query_proj(pooled_cls)                 # (B, n_tools_out)
+        # Dot product per candidate: (B, 1, d) * (B, K, d) sum→ (B, K)
+        scores = (tool_emb * query.unsqueeze(1)).sum(dim=-1)
+        temp = self.log_temperature.exp().clamp(max=10.0)
+        scores = scores * temp
+        return scores.masked_fill(~mask, -1e9)
+
+    @staticmethod
+    def contrastive_loss(
+        logits: torch.Tensor,
+        target_idx: torch.Tensor,
+        valid: torch.Tensor,
+        slot_0_weight: float = 1.0,
+    ) -> torch.Tensor:
+        """Cross-entropy on rows with a valid target. target_idx == -1
+        marks rows we should skip (no labeled positive in the
+        candidate set). Returns a scalar; returns zero if no row is
+        valid (so the optimiser sees no gradient from this batch).
+
+        ``slot_0_weight`` rescales the loss contribution for rows whose
+        target is slot-0 (the synthetic ``<no_tool>`` class). The cognos-
+        toolbox corpus is 63% ``<no_tool>``, which biases the head's
+        argmax toward slot-0. Setting ``slot_0_weight < 1.0`` re-balances
+        so real-tool slots get more gradient. ``1.0`` (default) preserves
+        the original unweighted behavior."""
+        if not bool(valid.any()):
+            return logits.sum() * 0.0   # zero-grad scalar with graph
+        weight = None
+        if slot_0_weight != 1.0:
+            K = logits.shape[-1]
+            weight = torch.ones(K, device=logits.device, dtype=logits.dtype)
+            weight[0] = float(slot_0_weight)
+        return torch.nn.functional.cross_entropy(
+            logits[valid], target_idx[valid].long(), weight=weight,
+        )