@10et/cli 1.0.0

package/scripts/train/v2/infer_server.py

@@ -0,0 +1,224 @@
+"""
+v2 Policy Head Inference Server — keeps model + cache warm in memory.
+
+Communicates via stdin/stdout JSONL (one request per line, one response per line).
+Eliminates ~3s model loading + ~0.1s cache loading per inference call.
+
+Start: python infer_server.py --checkpoint path/to/best_policy_head.pt
+Send:  {"state": "...", "goal": "...", "top_k": 3}
+Recv:  {"action": "fix_bug", "confidence": 0.94, "alternatives": [...]}
+
+Special commands:
+  {"cmd": "ping"}  → {"status": "ok", "cache_hits": N, "cache_misses": N}
+  {"cmd": "quit"}  → server exits
+"""
+
+import json
+import os
+import sys
+import time
+
+import torch
+import numpy as np
+
+# Add script directory to path for model import
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from model import PolicyHead
+
+
+class InferenceServer:
+    def __init__(self, checkpoint_path: str, cache_dir: str = None):
+        self.hits = 0
+        self.misses = 0
+        self.inferences = 0
+
+        # Device
+        if torch.cuda.is_available():
+            self.device = "cuda"
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            self.device = "mps"
+        else:
+            self.device = "cpu"
+
+        # Load model
+        t0 = time.time()
+        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+        config = ckpt["config"]
+        self.tool_to_index = ckpt["tool_to_index"]
+        self.index_to_tool = {v: k for k, v in self.tool_to_index.items()}
+
+        self.model = PolicyHead(
+            embedding_dim=config["embedding_dim"],
+            hidden_dim=config["hidden_dim"],
+            num_tools=ckpt["num_tools"],
+            num_layers=config["num_layers"],
+            num_heads=config["num_heads"],
+            dropout=config.get("dropout", 0.1),
+        ).to(self.device)
+        self.model.load_state_dict(ckpt["model_state_dict"])
+        self.model.eval()
+        self.embedding_dim = config["embedding_dim"]
+        t_model = time.time()
+
+        # Load embedding cache
+        self.text_to_idx = {}
+        self.embeddings = None
+        if cache_dir:
+            npz_path = os.path.join(cache_dir, "embeddings_cache.npz")
+            idx_path = os.path.join(cache_dir, "text_to_idx.json")
+            if os.path.exists(npz_path) and os.path.exists(idx_path):
+                data = np.load(npz_path)
+                self.embeddings = data["embeddings"]
+                self.text_to_idx = json.load(open(idx_path))
+        t_cache = time.time()
+
+        # In-memory cache for API-fetched embeddings (survives across requests)
+        self.live_cache: dict[str, list[float]] = {}
+
+        # Stratus API config
+        self.api_url = os.environ.get("STRATUS_API_URL", "https://api.stratus.run")
+        self.api_key = os.environ.get("STRATUS_API_KEY", "")
+
+        print(json.dumps({
+            "status": "ready",
+            "device": self.device,
+            "cache_size": len(self.text_to_idx),
+            "model_load_ms": int((t_model - t0) * 1000),
+            "cache_load_ms": int((t_cache - t_model) * 1000),
+        }), flush=True)
+
+    def get_embedding(self, text: str) -> list[float]:
+        """Get embedding: file cache → live cache → API → zero vector."""
+        # 1. File cache (precomputed)
+        if text in self.text_to_idx:
+            self.hits += 1
+            idx = self.text_to_idx[text]
+            return self.embeddings[idx].tolist()
+
+        # 2. Live cache (API results from this session)
+        if text in self.live_cache:
+            self.hits += 1
+            return self.live_cache[text]
+
+        # 3. API call
+        self.misses += 1
+        if self.api_key:
+            try:
+                import requests
+                response = requests.post(
+                    f"{self.api_url}/v1/embeddings",
+                    headers={
+                        "Authorization": f"Bearer {self.api_key}",
+                        "Content-Type": "application/json",
+                    },
+                    json={"model": "stratus-x1ac-base", "input": text},
+                    timeout=15,
+                )
+                response.raise_for_status()
+                emb = response.json()["data"][0]["embedding"]
+                self.live_cache[text] = emb  # Cache for future requests
+                return emb
+            except Exception as e:
+                print(json.dumps({"warning": f"API embed failed: {str(e)[:100]}"}), flush=True)
+
+        # 4. Zero vector fallback
+        return [0.0] * self.embedding_dim
+
+    def infer(self, state_text: str, goal_text: str, top_k: int = 3) -> dict:
+        t0 = time.time()
+
+        state_emb = self.get_embedding(state_text)
+        goal_emb = self.get_embedding(goal_text)
+        t_embed = time.time()
+
+        state_tensor = torch.tensor([state_emb], dtype=torch.float32).to(self.device)
+        goal_tensor = torch.tensor([goal_emb], dtype=torch.float32).to(self.device)
+
+        with torch.no_grad():
+            result = self.model.predict(state_tensor, goal_tensor, top_k=top_k)
+
+        top_indices = result["top_k_indices"][0].cpu().tolist()
+        top_probs = result["top_k_probs"][0].cpu().tolist()
+        t_infer = time.time()
+
+        self.inferences += 1
+
+        predictions = []
+        for idx, prob in zip(top_indices, top_probs):
+            predictions.append({
+                "action": self.index_to_tool[idx],
+                "confidence": round(prob, 4),
+            })
+
+        return {
+            "action": predictions[0]["action"],
+            "confidence": predictions[0]["confidence"],
+            "alternatives": predictions[1:],
+            "_timing_ms": {
+                "embed": int((t_embed - t0) * 1000),
+                "infer": int((t_infer - t_embed) * 1000),
+                "total": int((t_infer - t0) * 1000),
+            },
+        }
+
+    def run(self):
+        """Main loop: read JSONL from stdin, write JSONL to stdout."""
+        for line in sys.stdin:
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                req = json.loads(line)
+
+                # Handle special commands
+                if "cmd" in req:
+                    if req["cmd"] == "ping":
+                        print(json.dumps({
+                            "status": "ok",
+                            "inferences": self.inferences,
+                            "cache_hits": self.hits,
+                            "cache_misses": self.misses,
+                            "live_cache_size": len(self.live_cache),
+                        }), flush=True)
+                        continue
+                    elif req["cmd"] == "quit":
+                        break
+                    else:
+                        print(json.dumps({"error": f"Unknown command: {req['cmd']}"}), flush=True)
+                        continue
+
+                # Regular inference
+                state = req.get("state", "")
+                goal = req.get("goal", "")
+                top_k = req.get("top_k", 3)
+
+                result = self.infer(state, goal, top_k)
+                print(json.dumps(result), flush=True)
+
+            except Exception as e:
+                print(json.dumps({"error": str(e)}), flush=True)
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="v2 policy head inference server")
+    parser.add_argument("--checkpoint", required=True, help="Path to .pt checkpoint")
+    parser.add_argument("--cache-dir", default=None, help="Embedding cache directory")
+    args = parser.parse_args()
+
+    # Auto-find cache dir
+    if not args.cache_dir:
+        import pathlib
+        ckpt_dir = pathlib.Path(args.checkpoint).parent
+        for candidate in [ckpt_dir.parent / "v2-data", pathlib.Path.cwd() / ".jfl" / "v2-data"]:
+            if (candidate / "embeddings_cache.npz").exists():
+                args.cache_dir = str(candidate)
+                break
+
+    server = InferenceServer(args.checkpoint, args.cache_dir)
+    server.run()
+
+
+if __name__ == "__main__":
+    main()

package/scripts/train/v2/model.py

@@ -0,0 +1,112 @@
+"""
+v2 Policy Head — Transformer-based action selector.
+
+Architecture from Drew's Stratus tutorial:
+  (current_state_emb, goal_emb) -> state_proj + goal_proj -> fusion -> TransformerEncoder -> classifier -> action logits
+
+~8.7M params, ~17MB checkpoint. Replaces v1 MLP reward predictor.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class PolicyHead(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int = 768,
+        hidden_dim: int = 512,
+        num_tools: int = 12,
+        num_layers: int = 4,
+        num_heads: int = 8,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+
+        self.embedding_dim = embedding_dim
+        self.hidden_dim = hidden_dim
+        self.num_tools = num_tools
+
+        self.state_proj = nn.Linear(embedding_dim, hidden_dim)
+        self.goal_proj = nn.Linear(embedding_dim, hidden_dim)
+
+        self.fusion = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+        )
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dim,
+            nhead=num_heads,
+            dim_feedforward=hidden_dim * 4,
+            dropout=dropout,
+            activation="gelu",
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=num_layers,
+        )
+
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.classifier = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_tools),
+        )
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        current_state_emb: torch.Tensor,
+        goal_state_emb: torch.Tensor,
+    ) -> torch.Tensor:
+        state_h = self.state_proj(current_state_emb)
+        goal_h = self.goal_proj(goal_state_emb)
+
+        fused = self.fusion(torch.cat([state_h, goal_h], dim=-1))
+
+        x = fused.unsqueeze(1)
+        x = self.transformer(x)
+        x = x.squeeze(1)
+
+        x = self.norm(x)
+        logits = self.classifier(x)
+
+        return logits
+
+    def predict(
+        self,
+        current_state_emb: torch.Tensor,
+        goal_state_emb: torch.Tensor,
+        top_k: int = 3,
+    ) -> dict:
+        self.eval()
+        with torch.no_grad():
+            logits = self.forward(current_state_emb, goal_state_emb)
+            probs = F.softmax(logits, dim=-1)
+            top_probs, top_indices = torch.topk(probs, k=min(top_k, self.num_tools), dim=-1)
+
+        return {
+            "top_k_indices": top_indices,
+            "top_k_probs": top_probs,
+            "all_probs": probs,
+        }
+
+    @property
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)