caudate-cli 0.1.0__py3-none-any.whl

nn/caudate.py

@@ -0,0 +1,136 @@
+"""Caudate — Cognos's neural-network brain.
+
+A 4.6M-parameter transformer with text/vision/tool embedders, four
+attention layers, and a registry of output heads. The trunk is shared;
+each head specialises one routing job (which tool, which tier, should-
+think, expected-reward, plus growth slots for memory-write, cache,
+permission). Adding a new routing job means appending one HeadSpec to
+the registry — see `nn/heads.py`. Caudate is *not* a router; routing
+is the family of jobs the brain currently does.
+
+Stack:
+    StateEncoder ─► TransformerEncoder (4 layers, pre-norm) ─► CLS pool
+                                                              │
+                                                       + source_embed bias
+                                                              │
+                                                       HeadRegistry → {head_name: tensor}
+
+PyTorch's stock TransformerEncoder is used so the architecture stays
+inspectable: multi-head self-attention, feed-forward, residual,
+LayerNorm. Xavier init on the heads so the early loss isn't dominated
+by random output scale.
+
+Caudate runs as an *advisor* alongside the agent — predictions are
+logged but don't override the LLM's choices until trust is established.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from nn.config import NNConfig
+from nn.encoder import StateEncoder
+from nn.heads import ALL_HEADS, ContrastiveToolHead, HeadRegistry, HeadSpec
+
+
+class Caudate(nn.Module):
+    """End-to-end neural network. Forward: state → registered head outputs."""
+
+    name: str = "Caudate"
+
+    def __init__(self, cfg: NNConfig, head_specs: tuple[HeadSpec, ...] | None = None):
+        super().__init__()
+        self.cfg = cfg
+        self.encoder = StateEncoder(cfg)
+
+        layer = nn.TransformerEncoderLayer(
+            d_model=cfg.d_model,
+            nhead=cfg.n_heads,
+            dim_feedforward=cfg.d_ff,
+            dropout=cfg.dropout,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,           # pre-norm — more stable on small data
+        )
+        self.transformer = nn.TransformerEncoder(layer, num_layers=cfg.n_layers)
+        self.norm = nn.LayerNorm(cfg.d_model)
+
+        # Standard heads (tier/think/value + extended D-heads). The
+        # tool head was removed from this registry in Phase 2 of the
+        # format migration and replaced with `self.tool_head`, a
+        # contrastive open-vocab predictor that takes per-sample
+        # candidate tool definitions. The registry pattern still
+        # drives every OTHER head (uniform loss path in the trainer).
+        if head_specs is None:
+            specs = ALL_HEADS
+        else:
+            specs = head_specs
+        self.heads = HeadRegistry(cfg.d_model, cfg.dropout, specs)
+
+        # Phase-2 open-vocab tool predictor. Shares the state encoder's
+        # text embedder so we don't double-load sentence-transformers.
+        self.tool_head = ContrastiveToolHead(cfg, self.encoder.text_embedder)
+
+        # Phase 2 of CAUDATE_EVOLUTION.md: source-conditioned predictions.
+        # A learned per-teacher-model bias added to the pooled CLS token
+        # before the heads. Initialized to zero so legacy checkpoints
+        # loaded with strict=False (no source_embed in state_dict) keep
+        # their original behaviour until the next retrain.
+        n_sources = max(1, int(getattr(cfg, "source_vocab_size", 16)))
+        self.source_embed = nn.Embedding(n_sources, cfg.d_model)
+        nn.init.zeros_(self.source_embed.weight)
+
+    def forward(
+        self,
+        messages: list[list[str]],
+        tool_ids: torch.Tensor,
+        mood: torch.Tensor,
+        image_paths: list[list[str]] | None = None,
+        source_id: torch.Tensor | None = None,
+        tool_specs: list[list[str]] | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """tool_specs: per-sample list of "{name}: {description}"
+        strings. When supplied, the contrastive tool head runs and
+        emits `tool_logits` + `tool_mask` in the output dict. When
+        omitted or empty, those keys are absent (the trainer/advisor
+        treats that as "no tool prediction available for this sample").
+        """
+        x = self.encoder(messages, tool_ids, mood, image_paths)  # (B, L, d)
+        h = self.transformer(x)                           # (B, L, d)
+        cls = self.norm(h[:, 0])                          # (B, d)
+
+        # Source-conditioning: add a learned bias keyed on which model
+        # produced this turn's response. Default to slot 0 (<unknown>)
+        # which is initialized to zero — i.e. unconditioned baseline.
+        if source_id is None:
+            source_id = torch.zeros(cls.shape[0], dtype=torch.long, device=cls.device)
+        else:
+            source_id = source_id.to(device=cls.device, dtype=torch.long)
+            # Clamp out-of-range ids back to <unknown> rather than crash
+            n_sources = self.source_embed.num_embeddings
+            source_id = source_id.clamp(min=0, max=n_sources - 1)
+        cls = cls + self.source_embed(source_id)
+
+        head_out = self.heads(cls)            # dict of {output_key: tensor}
+        head_out["pooled"] = cls              # (B, d) — for inspection
+
+        # Contrastive tool head — runs only when candidates are
+        # supplied. Empty/missing → no tool prediction for this batch.
+        if tool_specs is not None and any(ts for ts in tool_specs):
+            tool_emb, tool_mask = self.tool_head.embed_tools(tool_specs)
+            head_out["tool_logits"] = self.tool_head(cls, tool_emb, tool_mask)
+            head_out["tool_mask"] = tool_mask
+
+        return head_out
+
+    @property
+    def head_specs(self) -> tuple[HeadSpec, ...]:
+        return self.heads.specs
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+
+
+# Backward-compat alias — old code that imported CognosController still works.
+CognosController = Caudate

nn/config.py

@@ -0,0 +1,141 @@
+"""Neural-network hyperparameters.
+
+Single source of truth for the controller's shape, training schedule,
+and runtime behavior. Persisted alongside checkpoints so loaded weights
+always match the matching architecture.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+
+class NNConfig(BaseModel):
+    """Hyperparameters + paths for the cognitive controller."""
+
+    # ----- Architecture -----
+    text_embed_dim: int = 384         # sentence-transformers MiniLM dim
+    # Phase-2: bumped 64 → 4096 to accommodate public tool-use corpora
+    # (xLAM has ~10K unique tool names; ToolBench similar). Tool ids
+    # that still overflow this cap are clamped at collate time so the
+    # embedding lookup never crashes — the cap is a hard ceiling.
+    tool_vocab_size: int = 4096       # max distinct tools we'll embed
+    tool_embed_dim: int = 64
+    # Phase 2 of CAUDATE_EVOLUTION.md: source-conditioned predictions.
+    # Each ConversationSample is tagged with the model that produced it;
+    # we look that name up in a small SourceVocab and bias the encoder's
+    # pooled state with a learned per-source embedding. Slot 0 is
+    # reserved for "<unknown>" so legacy untagged samples (and any
+    # model id we haven't seen before) get the zero baseline.
+    source_vocab_size: int = 16       # cap on distinct teacher-model ids
+    mood_dim: int = 4                 # confidence, curiosity, frustration, satisfaction
+    history_window: int = 16          # last N tool calls fed to encoder
+    # Phase-2 format migration: bumped 4 → 12 so multi-turn ToolBench-
+    # /xLAM-style conversations fit in the encoder window. Pos-embed
+    # shape changes with this, so any pre-Phase-2 checkpoint is
+    # incompatible — fresh retrain required.
+    msg_window: int = 12              # last N messages fed to encoder
+
+    # Vision channel — frozen image embedder.
+    use_vision: bool = True
+    image_window: int = 4             # last N images fed to encoder
+    # Backend = "clip" (light, 512-D) or "internvl" (rich, 4096-D).
+    # Match `vision_embed_dim` to whichever you pick.
+    vision_backend: str = "internvl"
+    vision_embed_dim: int = 4096
+    vision_encoder_name: str = "OpenGVLab/InternVL2-8B"
+    # InternVL only: the dtype to load with. bfloat16 needs ~16GB
+    # for InternVL2-8B on bf16-capable GPUs; float16 is similar; if
+    # you only have CPU set this to "float32".
+    vision_dtype: str = "bfloat16"
+
+    d_model: int = 256                # transformer hidden size
+    n_heads: int = 8
+    n_layers: int = 4
+    d_ff: int = 1024
+    dropout: float = 0.1
+
+    # ----- Output heads -----
+    # Phase-2 tool head is contrastive (open-vocab) — at inference the
+    # set of "tools the assistant can call right now" is supplied, and
+    # the model scores each candidate against the pooled CLS via a
+    # learned projection. n_tools_out is the projection's output
+    # dimension (= tool description embedding dim). 256 is roomy
+    # enough to cleanly separate the function-calling tool vocab in
+    # ToolBench / xLAM (~10K unique names) without ballooning params.
+    n_tools_out: int = 256
+    # Max tools considered per turn. Datasets typically expose 1-20;
+    # we pad to this cap with zero embeddings + a mask so the batch
+    # tensor shape is fixed.
+    max_tools_per_sample: int = 32
+    # value head outputs a single scalar; tier head outputs 2 logits;
+    # think head outputs 1 logit (sigmoid).
+
+    # ----- Training -----
+    batch_size: int = 32
+    learning_rate: float = 3e-4
+    weight_decay: float = 1e-2
+    grad_clip: float = 1.0
+    max_steps: int = 5000
+    eval_every: int = 200
+    save_every: int = 500
+    eval_split: float = 0.15
+    # Cap eval batches during training so the intra-fit eval ticks
+    # stay cheap on big public corpora. 100 batches × default 16 =
+    # 1600 samples — enough for a stable loss estimate, fast enough
+    # to not dominate the training loop. The end-of-fit final eval
+    # ignores this cap and runs over the full held-out split.
+    eval_max_batches_intra: int = 100
+    seed: int = 42
+
+    # Loss weights — controller has 4 heads, balance them so no one
+    # signal dominates training.
+    w_tool: float = 1.0
+    w_tier: float = 0.5
+    w_think: float = 0.2
+    w_value: float = 0.5
+
+    # Class weight on the contrastive tool head's slot-0 (<no_tool>).
+    # Cognos-toolbox is ~63% <no_tool>, which biased R3 step-11000 to
+    # argmax slot-0 on 100% of real-tool turns (target-in-top-5 was
+    # 52.5%, but slot-0 always won). 1.0 = unweighted (original
+    # behavior). 0.5 = moderate downweight; 0.3 = aggressive.
+    tool_no_tool_class_weight: float = 1.0
+
+    # ----- Replay buffer -----
+    replay_capacity: int = 8192
+    min_episodes_to_train: int = 64
+
+    # ----- Runtime -----
+    advisor_min_confidence: float = 0.55  # below this, advisor stays silent
+    # ADR-0007 constrained routing — at ADVISOR+, when tool_confidence
+    # clears `advisor_min_confidence` and the prediction is a real tool
+    # (not <no_tool>), the LLM's tool list is pruned to this many of
+    # Caudate's top picks. The predicted tool itself is always included,
+    # and `advisor_top_k_floor` is enforced so the LLM never sees fewer
+    # than that many options (mitigates "confident wrong hides the right
+    # tool" per ADR-0007 Consequences).
+    advisor_top_k: int = 5
+    advisor_top_k_floor: int = 3
+    advisor_log_path: str = "data/nn/predictions.jsonl"
+
+    # ----- Paths -----
+    data_dir: str = "data/nn"
+    checkpoint_path: str = "data/nn/caudate.pt"
+    metadata_path: str = "data/nn/caudate.meta.json"
+    text_encoder_name: str = "sentence-transformers/all-MiniLM-L6-v2"
+
+    # Special tokens for the tool history channel
+    tool_pad_token: int = 0
+    tool_unk_token: int = 1
+    tool_bos_token: int = 2
+
+    @property
+    def total_input_dim(self) -> int:
+        return self.text_embed_dim + self.tool_embed_dim + self.mood_dim
+
+    def ensure_dirs(self) -> None:
+        Path(self.data_dir).mkdir(parents=True, exist_ok=True)
+        Path(self.advisor_log_path).parent.mkdir(parents=True, exist_ok=True)

nn/consolidator.py

@@ -0,0 +1,81 @@
+"""Memory consolidator — periodic offline replay → training pairs.
+
+Loosely modeled on hippocampal-cortical consolidation: while Cognos is
+idle (or on demand), it walks recent episodic memories, distills them
+into ConversationSample rows, and pushes them into the replay buffer.
+The trainer can then run additional gradient steps without needing
+fresh data from a live session.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Any
+
+from nn.config import NNConfig
+from nn.data import ConversationSample, ReplayBuffer
+from nn.format import ChatMessage, ToolCall
+
+logger = logging.getLogger(__name__)
+
+
+def consolidate_from_episodic(
+    episodic: Any,
+    replay: ReplayBuffer,
+    cfg: NNConfig,
+    limit: int = 128,
+) -> int:
+    """Pull up to `limit` episodes from episodic memory into the replay buffer.
+
+    Returns number of samples added. If episodic memory exposes a
+    `recall_recent(limit)` method we use that; otherwise we degrade to
+    `recall("", limit=...)` which most stores accept.
+    """
+    if episodic is None:
+        return 0
+    try:
+        if hasattr(episodic, "recall_recent"):
+            episodes = episodic.recall_recent(limit=limit)
+        else:
+            episodes = episodic.recall("", limit=limit)
+    except Exception as e:
+        logger.warning(f"consolidator: episodic recall failed: {e}")
+        return 0
+
+    added = 0
+    for ep in episodes or []:
+        try:
+            tool = getattr(ep, "tool_name", None) or "unknown"
+            args = getattr(ep, "tool_args", {}) or {}
+            result = getattr(ep, "result", None)
+            success = (
+                getattr(result, "status", None)
+                and str(getattr(result, "status").value if hasattr(getattr(result, "status"), "value") else getattr(result, "status")) == "success"
+            )
+            action_text = (getattr(ep, "action", "") or "")[:400]
+            try:
+                args_str = json.dumps(args) if args else ""
+            except Exception:
+                args_str = ""
+            sample = ConversationSample(
+                # One episodic memory → a 2-turn conversation: the
+                # action description as a user turn, followed by the
+                # assistant calling the tool that was actually invoked.
+                conversation=[
+                    ChatMessage(role="user", content=action_text),
+                    ChatMessage(
+                        role="assistant", content="",
+                        tool_calls=[ToolCall(name=tool, arguments=args_str)],
+                    ),
+                ],
+                tools=[],
+                mood=[0.5, 0.5, 0.5, 0.5],
+                target_tool=tool,
+                target_value=(0.7 if success else 0.3),
+            )
+            replay.push(sample)
+            added += 1
+        except Exception as e:
+            logger.debug(f"consolidator: skipped episode: {e}")
+    return added