caudate-cli 0.1.0__py3-none-any.whl

nn/encoder.py

@@ -0,0 +1,258 @@
+"""State encoder — turns Cognos's per-turn state into a tensor.
+
+Three channels are fused into a sequence the controller can attend over:
+
+    1. Text channel — recent N messages embedded with sentence-transformers.
+       Frozen during training (kept lightweight) — we learn on top of it
+       rather than fine-tuning the embedder.
+
+    2. Tool history channel — the last K tool calls as integer ids,
+       embedded with a learned `nn.Embedding`.
+
+    3. Mood channel — 4 continuous floats projected into d_model space.
+
+The output of the encoder is a (B, L, d_model) tensor that flows into the
+transformer controller. L = msg_window + history_window + 1 (mood token).
+A single learned [CLS] token is prepended so heads can read off a fixed
+position regardless of how the channels are weighted at attention time.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any
+
+# Defensive — block transformers' TF import path before sentence-transformers loads.
+os.environ.setdefault("USE_TF", "0")
+
+import torch
+import torch.nn as nn
+
+from nn.config import NNConfig
+from nn.vision import VisionEncoder
+
+logger = logging.getLogger(__name__)
+
+
+class _SentenceEmbedder:
+    """Lazy wrapper around sentence-transformers. Falls back to a deterministic
+    hash-based fake embedder if the lib isn't installed — so the rest of the
+    pipeline (data loading, training loop wiring) still exercises end-to-end
+    in an offline test environment."""
+
+    def __init__(self, model_name: str, embed_dim: int):
+        self.model_name = model_name
+        self.embed_dim = embed_dim
+        self._model = None
+        self._fake = False
+
+    def _load(self) -> None:
+        if self._model is not None or self._fake:
+            return
+        try:
+            from sentence_transformers import SentenceTransformer
+            # Place the embedder on CUDA when available. Without this the
+            # MiniLM forward runs on CPU and the numpy round-trip dominates
+            # step time (was ~3 s/step with a 5.5M trunk on a 3090).
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            self._model = SentenceTransformer(self.model_name, device=device)
+            self._device = device
+            real_dim = self._model.get_sentence_embedding_dimension()
+            if real_dim != self.embed_dim:
+                logger.warning(
+                    f"Embedder dim {real_dim} != configured {self.embed_dim} — "
+                    "you'll need to retrain if you change models"
+                )
+        except Exception as e:
+            logger.warning(
+                f"sentence-transformers unavailable ({e}); "
+                "falling back to deterministic hash embedder"
+            )
+            self._fake = True
+
+    def encode(self, texts: list[str]) -> torch.Tensor:
+        self._load()
+        if self._fake or self._model is None:
+            return _hash_embed(texts, self.embed_dim)
+        with torch.no_grad():
+            # Keep the output as a tensor on the embedder's device — no
+            # CPU/numpy round-trip. Caller (.to(self.device) in encoder.py)
+            # handles any final placement if devices differ.
+            t = self._model.encode(
+                texts,
+                convert_to_tensor=True,
+                normalize_embeddings=True,
+            )
+        # sentence-transformers runs encode() under torch.inference_mode().
+        # The returned tensor is an "inference tensor" that autograd refuses
+        # to use downstream (the old numpy round-trip detached this flag
+        # accidentally). Clone strips it.
+        return t.clone().float()
+
+
+def _hash_embed(texts: list[str], dim: int) -> torch.Tensor:
+    """Deterministic per-character hash → fixed-dim vector. NOT semantically
+    meaningful; only useful as a placeholder so the pipeline runs offline."""
+    out = torch.zeros((len(texts), dim), dtype=torch.float32)
+    for i, t in enumerate(texts):
+        for j, ch in enumerate((t or "")[: dim * 2]):
+            out[i, (j * 16807 + ord(ch)) % dim] += (ord(ch) % 17) / 17.0
+    norm = out.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+    return out / norm
+
+
+class StateEncoder(nn.Module):
+    """Encode (messages, tool_history, mood, images) → (B, L, d_model)."""
+
+    def __init__(self, cfg: NNConfig):
+        super().__init__()
+        self.cfg = cfg
+
+        # Sentence-transformer is frozen — we project its output instead.
+        self.text_embedder = _SentenceEmbedder(cfg.text_encoder_name, cfg.text_embed_dim)
+        self.text_proj = nn.Linear(cfg.text_embed_dim, cfg.d_model)
+
+        # Learned tool-id embedding for the history channel.
+        self.tool_embed = nn.Embedding(
+            cfg.tool_vocab_size, cfg.tool_embed_dim,
+            padding_idx=cfg.tool_pad_token,
+        )
+        self.tool_proj = nn.Linear(cfg.tool_embed_dim, cfg.d_model)
+
+        # Mood is 4 floats — project up to d_model.
+        self.mood_proj = nn.Sequential(
+            nn.Linear(cfg.mood_dim, cfg.d_model),
+            nn.GELU(),
+            nn.Linear(cfg.d_model, cfg.d_model),
+        )
+
+        # Vision channel — frozen image embedder, projected to d_model.
+        # Backend is chosen at config time (CLIP for light, InternVL2
+        # for rich semantic features). Caudate learns on top of these.
+        if cfg.use_vision:
+            from nn.vision import make_vision_encoder
+            self.vision_encoder = make_vision_encoder(
+                backend=cfg.vision_backend,
+                model_name=cfg.vision_encoder_name,
+                dtype=getattr(cfg, "vision_dtype", "bfloat16"),
+            )
+            self.vision_proj = nn.Sequential(
+                nn.Linear(cfg.vision_embed_dim, cfg.d_model),
+                nn.GELU(),
+                nn.Linear(cfg.d_model, cfg.d_model),
+            )
+        else:
+            self.vision_encoder = None
+            self.vision_proj = None
+
+        # Channel-type embedding so the transformer knows which kind of
+        # token it's looking at (cls=0, text=1, tool=2, mood=3, image=4).
+        n_channel_types = 5 if cfg.use_vision else 4
+        self.type_embed = nn.Embedding(n_channel_types, cfg.d_model)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.d_model))
+        nn.init.normal_(self.cls_token, std=0.02)
+
+        # Position embedding across the fused sequence.
+        max_len = (
+            1 + cfg.msg_window + cfg.history_window + 1
+            + (cfg.image_window if cfg.use_vision else 0)
+        )
+        self.pos_embed = nn.Embedding(max_len, cfg.d_model)
+        self.dropout = nn.Dropout(cfg.dropout)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def encode_messages(self, messages: list[list[str]]) -> torch.Tensor:
+        """Embed a batch of message lists. messages[b] is a list of strings.
+
+        Pads / truncates to msg_window on the right.
+        """
+        B = len(messages)
+        W = self.cfg.msg_window
+        padded: list[list[str]] = []
+        for batch in messages:
+            batch = list(batch)[-W:]
+            if len(batch) < W:
+                batch = [""] * (W - len(batch)) + batch
+            padded.append(batch)
+
+        flat = [m for batch in padded for m in batch]
+        emb = self.text_embedder.encode(flat).to(self.device)  # (B*W, text_embed_dim)
+        emb = emb.view(B, W, self.cfg.text_embed_dim)
+        return emb
+
+    def encode_images(
+        self, image_paths: list[list[str]],
+    ) -> torch.Tensor:
+        """Embed a batch of per-sample image-path lists.
+
+        image_paths[b] is a list of paths to images attached at this
+        turn. Pads / truncates to image_window on the right. Returns
+        (B, image_window, vision_embed_dim) float.
+        """
+        if not self.cfg.use_vision or self.vision_encoder is None:
+            return torch.zeros(0)
+        B = len(image_paths)
+        W = self.cfg.image_window
+        flat: list[str] = []
+        for paths in image_paths:
+            paths = list(paths)[-W:]
+            if len(paths) < W:
+                paths = [""] * (W - len(paths)) + paths
+            flat.extend(paths)
+        emb = self.vision_encoder.encode_paths(flat).to(self.device)
+        return emb.view(B, W, self.cfg.vision_embed_dim)
+
+    def forward(
+        self,
+        messages: list[list[str]],
+        tool_ids: torch.Tensor,         # (B, history_window) long
+        mood: torch.Tensor,             # (B, mood_dim) float
+        image_paths: list[list[str]] | None = None,  # (B,) list of paths
+    ) -> torch.Tensor:
+        B = len(messages)
+        W_msg = self.cfg.msg_window
+        W_tool = self.cfg.history_window
+        W_img = self.cfg.image_window if self.cfg.use_vision else 0
+
+        # 1. text channel
+        text_emb = self.encode_messages(messages)
+        text_proj = self.text_proj(text_emb)                            # (B, W_msg, d)
+
+        # 2. tool channel
+        tool_emb = self.tool_embed(tool_ids.to(self.device))            # (B, W_tool, te)
+        tool_proj = self.tool_proj(tool_emb)                            # (B, W_tool, d)
+
+        # 3. mood channel — single token
+        mood_proj = self.mood_proj(mood.to(self.device)).unsqueeze(1)   # (B, 1, d)
+
+        # 4. vision channel (optional)
+        cls = self.cls_token.expand(B, -1, -1)
+        chunks = [cls, text_proj, tool_proj, mood_proj]
+        if self.cfg.use_vision:
+            paths_in = image_paths if image_paths is not None else [[] for _ in range(B)]
+            img_emb = self.encode_images(paths_in)                       # (B, W_img, ve)
+            img_proj = self.vision_proj(img_emb)                         # (B, W_img, d)
+            chunks.append(img_proj)
+        x = torch.cat(chunks, dim=1)                                     # (B, L, d)
+
+        # 5. add type + positional embeddings
+        L = x.size(1)
+        type_ids = torch.zeros(L, dtype=torch.long, device=self.device)
+        # 0=cls, 1=text, 2=tool, 3=mood, 4=image
+        type_ids[1:1 + W_msg] = 1
+        type_ids[1 + W_msg:1 + W_msg + W_tool] = 2
+        type_ids[1 + W_msg + W_tool] = 3                                 # mood (1 token)
+        if self.cfg.use_vision:
+            start = 2 + W_msg + W_tool
+            type_ids[start:start + W_img] = 4
+        type_embedding = self.type_embed(type_ids).unsqueeze(0)          # (1, L, d)
+
+        pos_ids = torch.arange(L, device=self.device)
+        pos_embedding = self.pos_embed(pos_ids).unsqueeze(0)             # (1, L, d)
+
+        x = x + type_embedding + pos_embedding
+        return self.dropout(x)

nn/forge_advisor.py

@@ -0,0 +1,303 @@
+"""Forge feature-level advisor — Caudate's two-way loop with the
+orchestrator.
+
+ADR 0006 calls for a single predictor `predict_feature_difficulty`
+that drives four orchestrator behaviours:
+  B. Model selection per feature (system1 vs system2)
+  C. Backlog ordering when priorities tie
+  D. Early revision when a retry looks doomed
+
+This module exports the prediction surface and falls back to a
+**heuristic baseline** when no trained Caudate head is available yet.
+The baseline is computed from the rolling history in
+``data/nn/feature_outcomes.jsonl``:
+
+  - per-model success rate (global)
+  - per-model average n_turns on success vs failure
+  - keyword bias on the feature text (the agent succeeded on these
+    kinds of feature before / didn't)
+
+A swap to a real trained head changes only the body of
+``_predict_via_caudate``; the public functions and orchestrator
+call-sites stay identical.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import math
+import re
+import time
+from collections import Counter, defaultdict
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+_OUTCOMES_PATH = Path("data/nn/feature_outcomes.jsonl")
+_AUDIT_PATH = Path("data/nn/forge_predictions.jsonl")
+
+# Cache the loaded advisor so we don't reload Caudate's weights on
+# every feature decision. Cleared by reset_advisor_cache() in tests.
+_advisor_cache: object | None = None
+_advisor_cache_attempted: bool = False
+
+# Heuristic-baseline knobs. Reasonable defaults; tuned later from
+# real data when Caudate's head is trained.
+_DEFAULT_SUCCESS_PROB = 0.55
+_MAX_FEATURES_FOR_BASELINE = 2000  # bound the rolling window read
+
+
+@dataclass
+class FeaturePrediction:
+    """Result of ``predict_feature_difficulty``."""
+    success_prob: float          # [0, 1]
+    difficulty: float            # 1 - success_prob (convenience)
+    source: str                  # 'caudate' | 'baseline' | 'fallback'
+    reason: str                  # short human-readable explanation
+
+
+# ─── Public surface ───────────────────────────────────────────────────
+
+
+def predict_feature_difficulty(
+    feature_text: str,
+    model_id: str | None = None,
+) -> FeaturePrediction:
+    """Best-effort P(success | feature, model). Used by the orchestrator
+    to decide whether to promote a feature to system2, order the
+    backlog by predicted difficulty, or trigger early revision."""
+    pred = _predict_via_caudate(feature_text, model_id)
+    if pred is not None:
+        _audit(feature_text, model_id, pred)
+        return pred
+    pred = _predict_via_baseline(feature_text, model_id)
+    _audit(feature_text, model_id, pred)
+    return pred
+
+
+def predict_failure_repeats(
+    feature_text: str,
+    n_prior_failures: int,
+    model_id: str | None = None,
+) -> float:
+    """P(next retry also fails | this feature has already failed
+    ``n_prior_failures`` times). Used by ADR-0006 Phase D to decide
+    whether to skip straight to revision rather than burn more
+    sessions."""
+    base = predict_feature_difficulty(feature_text, model_id)
+    # Conditional on already failing N times, P(next fail) rises
+    # quickly. A simple Bayesian-flavoured update — each prior failure
+    # halves the remaining success probability.
+    success_left = base.success_prob * (0.5 ** max(0, n_prior_failures))
+    return float(min(0.99, 1.0 - success_left))
+
+
+# ─── Caudate path (real head — TODO: wire when trained) ──────────────
+
+
+def reset_advisor_cache() -> None:
+    """Drop the cached CaudateAdvisor — used by tests that swap weights."""
+    global _advisor_cache, _advisor_cache_attempted
+    _advisor_cache = None
+    _advisor_cache_attempted = False
+
+
+def _get_advisor():
+    """Load the CaudateAdvisor once and cache it. None means we tried
+    and the checkpoint isn't available — don't re-try every call."""
+    global _advisor_cache, _advisor_cache_attempted
+    if _advisor_cache_attempted:
+        return _advisor_cache
+    _advisor_cache_attempted = True
+    try:
+        from nn.runtime import load_advisor  # type: ignore
+    except Exception as e:
+        logger.debug(f"nn.runtime unavailable: {e}")
+        return None
+    try:
+        _advisor_cache = load_advisor()
+    except Exception as e:
+        logger.debug(f"load_advisor failed: {e}")
+        _advisor_cache = None
+    return _advisor_cache
+
+
+def _predict_via_caudate(
+    feature_text: str, model_id: str | None,
+) -> FeaturePrediction | None:
+    """Trained Caudate path. Returns None when:
+      - the checkpoint is missing,
+      - the feature_success head exists but is at its init weights
+        (predict_feature_success itself returns None in that case),
+      - inference raises.
+
+    The advisor instance is cached so we don't reload weights every
+    feature decision."""
+    advisor = _get_advisor()
+    if advisor is None:
+        return None
+    head_fn = getattr(advisor, "predict_feature_success", None)
+    if head_fn is None:
+        return None
+    try:
+        result = head_fn(feature_text=feature_text, model_id=model_id)
+    except Exception as e:
+        logger.debug(f"caudate feature head call raised: {e}")
+        return None
+    if not isinstance(result, dict) or "success_prob" not in result:
+        return None
+    p = float(max(0.0, min(1.0, result["success_prob"])))
+    return FeaturePrediction(
+        success_prob=p, difficulty=1 - p,
+        source="caudate",
+        reason=str(result.get("reason", "trained head")),
+    )
+
+
+# ─── Heuristic baseline ───────────────────────────────────────────────
+
+
+def _predict_via_baseline(
+    feature_text: str, model_id: str | None,
+) -> FeaturePrediction:
+    """Fall back to a corpus-driven heuristic. Reads feature_outcomes
+    JSONL, computes per-model success rate + keyword-weighted prior
+    on the feature text. Defaults to ``_DEFAULT_SUCCESS_PROB`` when no
+    history exists."""
+    history = _load_recent_outcomes()
+    if not history:
+        return FeaturePrediction(
+            success_prob=_DEFAULT_SUCCESS_PROB,
+            difficulty=1 - _DEFAULT_SUCCESS_PROB,
+            source="fallback",
+            reason="no feature-outcome history yet",
+        )
+
+    # Per-model success rate
+    by_model: dict[str, list[bool]] = defaultdict(list)
+    for h in history:
+        by_model[h.get("model_used", "unknown")].append(bool(h.get("success")))
+    model_rates = {
+        m: (sum(v) / len(v)) for m, v in by_model.items() if v
+    }
+
+    # Pick the rate that matches the requested model. Fall back to
+    # the global success rate if the model has no history.
+    rate = None
+    reason_parts: list[str] = []
+    if model_id and model_id in model_rates and len(by_model[model_id]) >= 3:
+        rate = model_rates[model_id]
+        reason_parts.append(
+            f"model={model_id} n={len(by_model[model_id])} rate={rate:.2f}"
+        )
+    if rate is None:
+        flat = [h for v in by_model.values() for h in v]
+        rate = sum(flat) / max(1, len(flat))
+        reason_parts.append(f"global n={len(flat)} rate={rate:.2f}")
+
+    # Keyword bias on the feature text. Tokens that appear in
+    # successful features more often than failed → bonus; reverse → penalty.
+    bias = _keyword_bias(feature_text, history)
+    if bias != 0.0:
+        reason_parts.append(f"kw_bias={bias:+.3f}")
+
+    p = max(0.05, min(0.95, rate + bias))
+    return FeaturePrediction(
+        success_prob=p,
+        difficulty=1 - p,
+        source="baseline",
+        reason=" · ".join(reason_parts),
+    )
+
+
+def _load_recent_outcomes() -> list[dict[str, Any]]:
+    """Read the most recent feature_outcomes rows. Skips garbage lines."""
+    if not _OUTCOMES_PATH.exists():
+        return []
+    try:
+        with _OUTCOMES_PATH.open() as f:
+            lines = f.readlines()
+    except OSError:
+        return []
+    rows: list[dict[str, Any]] = []
+    for line in lines[-_MAX_FEATURES_FOR_BASELINE:]:
+        try:
+            rows.append(json.loads(line))
+        except Exception:
+            continue
+    return rows
+
+
+_WORD_RE = re.compile(r"[a-zA-Z][a-zA-Z0-9]{3,}")
+_STOPWORDS = frozenset({
+    "feature", "implement", "create", "build", "with", "from", "that",
+    "this", "they", "their", "have", "will", "should", "would", "could",
+    "must", "needs", "need", "want", "wants", "make", "made", "makes",
+    "into", "onto", "over", "under", "between", "where", "when", "what",
+    "which", "while", "during", "before", "after", "above", "below",
+    "page", "pages", "user", "users", "data", "code", "test", "tests",
+})
+
+
+def _keyword_bias(feature_text: str, history: list[dict]) -> float:
+    """Compute a small additive bias [-0.15, +0.15] from the feature
+    text's distinctive tokens. Tokens that appear disproportionately
+    in past successes shift the prediction up; failures shift it
+    down. Bounded to keep heuristic noise from dominating."""
+    text = (feature_text or "").lower()
+    tokens = {t for t in _WORD_RE.findall(text) if t not in _STOPWORDS}
+    if not tokens:
+        return 0.0
+
+    # Build per-token success/fail counts on history
+    success_counts: Counter = Counter()
+    failure_counts: Counter = Counter()
+    for h in history:
+        h_text = (h.get("feature_text") or "").lower()
+        h_tokens = {
+            t for t in _WORD_RE.findall(h_text) if t not in _STOPWORDS
+        }
+        target = success_counts if h.get("success") else failure_counts
+        for tok in h_tokens:
+            target[tok] += 1
+
+    score = 0.0
+    for tok in tokens:
+        s, f = success_counts.get(tok, 0), failure_counts.get(tok, 0)
+        if s + f < 3:
+            continue
+        score += (s - f) / (s + f)
+    # Squash to [-0.15, +0.15] via tanh
+    return 0.15 * math.tanh(score / max(1, len(tokens)))
+
+
+def _audit(
+    feature_text: str, model_id: str | None, pred: FeaturePrediction,
+) -> None:
+    """Append a row to data/nn/forge_predictions.jsonl so we can later
+    correlate predictions with actual outcomes."""
+    try:
+        _AUDIT_PATH.parent.mkdir(parents=True, exist_ok=True)
+        with _AUDIT_PATH.open("a", encoding="utf-8") as f:
+            f.write(json.dumps({
+                "ts": time.time(),
+                "feature_text": feature_text[:200],
+                "model_id": model_id,
+                "success_prob": pred.success_prob,
+                "source": pred.source,
+                "reason": pred.reason,
+            }) + "\n")
+    except Exception as e:
+        logger.debug(f"forge_advisor audit write failed: {e}")
+
+
+__all__ = [
+    "FeaturePrediction",
+    "predict_feature_difficulty",
+    "predict_failure_repeats",
+    "reset_advisor_cache",
+]