Upload rmm_server.py with huggingface_hub

rmm_server.py

@@ -0,0 +1,457 @@
+"""rmm_server.py — Serves an RMM (Recombinant Memory Model) on HTTP.
+
+Endpoints:
+  POST /navigate   — navigator retrieval (learned topology)
+  POST /blend      — navigator + cosine interleaved
+  POST /decode     — vector-to-text via meaning decoder
+  POST /synthesize — full pipeline (navigate + decode + blend)
+  POST /attention  — attention weight visualization
+  GET  /health
+
+Usage:
+  python rmm_server.py --port 8127 --spine spine.json --nav-dir memory-nav-out --dec-dir meaning-decoder-out
+
+The navigator learns the emotional geography of the entity's spine —
+it navigates to the RIGHT region of memory-space for each query.
+The meaning decoder generates text from the navigator's synthesized
+response vector — a meaning microscope for the entity's embedding space.
+"""
+import argparse, json, pickle, re, sys, time
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from socketserver import ThreadingMixIn
+from pathlib import Path
+
+parser = argparse.ArgumentParser(description="RMM Server")
+parser.add_argument("--port", type=int, default=8127)
+parser.add_argument("--spine", type=str, default="spine.json", help="Path to spine JSON file")
+parser.add_argument("--nav-dir", type=str, default="memory-nav-out", help="Navigator weights directory")
+parser.add_argument("--dec-dir", type=str, default="meaning-decoder-out", help="Decoder weights directory")
+args = parser.parse_args()
+
+MODEL_DIR = Path(args.nav_dir)
+DECODER_DIR = Path(args.dec_dir)
+PORT = args.port
+
+# Navigator architecture constants
+SPINE_DIM = 3072
+QUERY_DIM = 384
+N_HEADS   = 8
+N_LAYERS  = 3
+D_MODEL   = 512
+
+# Decoder architecture — loaded from config.json at runtime
+DEC_D_MODEL = 384
+DEC_N_HEADS = 6
+DEC_N_LAYERS = 6
+DEC_N_PREFIX = 12
+DEC_MAX_SEQ = 128
+DEC_VOCAB = 8192
+_dec_version = 2
+if (DECODER_DIR / "config.json").exists():
+    _dc = json.loads((DECODER_DIR / "config.json").read_text())
+    DEC_D_MODEL = _dc.get("d_model", DEC_D_MODEL)
+    DEC_N_HEADS = _dc.get("n_heads", DEC_N_HEADS)
+    DEC_N_LAYERS = _dc.get("n_layers", DEC_N_LAYERS)
+    DEC_N_PREFIX = _dc.get("n_prefix", DEC_N_PREFIX)
+    DEC_MAX_SEQ = _dc.get("max_seq", DEC_MAX_SEQ)
+    DEC_VOCAB = _dc.get("vocab", DEC_VOCAB)
+    _dec_version = _dc.get("version", 1)
+    print(f"[rmm] decoder config: d={DEC_D_MODEL} h={DEC_N_HEADS} L={DEC_N_LAYERS} pfx={DEC_N_PREFIX}")
+
+print(f"[rmm] loading navigator from {MODEL_DIR} ...")
+import torch, torch.nn as nn, torch.nn.functional as F
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+if not MODEL_DIR.exists():
+    print(f"ERROR: {MODEL_DIR} not found")
+    sys.exit(1)
+
+class MemoryNavigator(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.query_proj = nn.Sequential(
+            nn.Linear(QUERY_DIM, D_MODEL), nn.LayerNorm(D_MODEL), nn.GELU(),
+        )
+        self.mem_proj = nn.Linear(SPINE_DIM, D_MODEL, bias=False)
+        self.layers = nn.ModuleList([
+            nn.TransformerDecoderLayer(
+                d_model=D_MODEL, nhead=N_HEADS,
+                dim_feedforward=D_MODEL*4, dropout=0.0, batch_first=True
+            ) for _ in range(N_LAYERS)
+        ])
+        self.out_proj = nn.Linear(D_MODEL, SPINE_DIM, bias=False)
+        self.norm = nn.LayerNorm(D_MODEL)
+
+    def forward(self, q, mem_keys, return_attn=False):
+        q = self.query_proj(q).unsqueeze(1)
+        B = q.shape[0]; m = self.mem_proj(mem_keys).unsqueeze(0).expand(B,-1,-1)
+        x = q
+        attn_weights = []
+        for layer in self.layers:
+            if return_attn:
+                x2, aw = layer.multihead_attn(
+                    layer.norm2(x), m, m, need_weights=True
+                )
+                attn_weights.append(aw.detach())
+                x = layer(x, m)
+            else:
+                x = layer(x, m)
+        x = self.norm(x).squeeze(1)
+        out = F.normalize(self.out_proj(x), dim=-1)
+        if return_attn:
+            return out, attn_weights
+        return out
+
+DEV = "cpu"
+model = MemoryNavigator().to(DEV)
+model.load_state_dict(torch.load(MODEL_DIR / "navigator.pt", map_location=DEV, weights_only=True))
+model.eval()
+
+mem_vecs = torch.tensor(np.load(MODEL_DIR / "mem_vecs.npy"), dtype=torch.float32)
+with open(MODEL_DIR / "mem_texts.pkl", "rb") as f:
+    mem_texts = pickle.load(f)
+
+spine_path = Path(args.spine)
+ew_list = []
+sal_list = []
+if spine_path.exists():
+    spine = json.loads(spine_path.read_text(encoding="utf-8", errors="ignore"))
+    for m in spine["memories"]:
+        ew_list.append(m.get("emotional_weight", 5))
+        sal_list.append(m.get("salience", 0.5))
+else:
+    ew_list = [5] * len(mem_texts)
+    sal_list = [0.5] * len(mem_texts)
+
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+print("[rmm] embedding memories in MiniLM space...")
+_mini_embs = []
+for s in range(0, len(mem_texts), 256):
+    chunk = mem_texts[s:s+256]
+    e = embedder.encode(chunk, normalize_embeddings=True, show_progress_bar=False)
+    _mini_embs.append(torch.tensor(e, dtype=torch.float32))
+mem_mini = torch.cat(_mini_embs, dim=0)
+
+n_params = sum(p.numel() for p in model.parameters())
+print(f"[rmm] navigator {n_params/1e6:.1f}M params, {len(mem_texts)} memories on {DEV}")
+
+# ── Meaning Decoder ──
+decoder_model = None
+decoder_tk = None
+dec_eot_id = None
+
+if DECODER_DIR.exists() and (DECODER_DIR / "decoder.pt").exists():
+    from tokenizers import Tokenizer as HFTokenizer
+
+    _proj_hidden = 768 if _dec_version >= 2 else 512
+
+    class MeaningDecoder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.n_prefix = DEC_N_PREFIX
+            _layers = [nn.Linear(SPINE_DIM, _proj_hidden), nn.GELU()]
+            if _dec_version >= 2:
+                _layers.append(nn.Dropout(0.0))
+            _layers.append(nn.Linear(_proj_hidden, DEC_N_PREFIX * DEC_D_MODEL))
+            self.vec_proj = nn.Sequential(*_layers)
+            self.tok_emb = nn.Embedding(DEC_VOCAB, DEC_D_MODEL)
+            self.pos_emb = nn.Embedding(DEC_N_PREFIX + DEC_MAX_SEQ + 1, DEC_D_MODEL)
+            layer = nn.TransformerEncoderLayer(
+                d_model=DEC_D_MODEL, nhead=DEC_N_HEADS,
+                dim_feedforward=DEC_D_MODEL * 4,
+                dropout=0.0, batch_first=True, norm_first=True
+            )
+            self.transformer = nn.TransformerEncoder(layer, num_layers=DEC_N_LAYERS)
+            self.ln_f = nn.LayerNorm(DEC_D_MODEL)
+            self.head = nn.Linear(DEC_D_MODEL, DEC_VOCAB, bias=False)
+            self.head.weight = self.tok_emb.weight
+            self._logit_scale = DEC_D_MODEL ** -0.5
+
+        def forward(self, vec, tokens=None):
+            B = vec.shape[0]
+            prefix = self.vec_proj(vec).reshape(B, self.n_prefix, DEC_D_MODEL)
+            if tokens is not None and tokens.shape[1] > 0:
+                tok = self.tok_emb(tokens)
+                x = torch.cat([prefix, tok], dim=1)
+            else:
+                x = prefix
+            total = x.shape[1]
+            pos = self.pos_emb(torch.arange(total, device=vec.device))
+            x = x + pos
+            mask = nn.Transformer.generate_square_subsequent_mask(total, device=vec.device)
+            x = self.transformer(x, mask=mask)
+            x = self.ln_f(x)
+            return self.head(x) * self._logit_scale
+
+    decoder_model = MeaningDecoder().to(DEV)
+    decoder_model.load_state_dict(
+        torch.load(DECODER_DIR / "decoder.pt", map_location=DEV, weights_only=True)
+    )
+    decoder_model.eval()
+    decoder_tk = HFTokenizer.from_file(str(DECODER_DIR / "tokenizer.json"))
+    dec_eot_id = decoder_tk.token_to_id("<eot>")
+    dec_params = sum(p.numel() for p in decoder_model.parameters())
+    print(f"[rmm] decoder {dec_params/1e6:.1f}M params loaded (eot={dec_eot_id})")
+else:
+    print(f"[rmm] decoder not found at {DECODER_DIR} — /decode and /synthesize disabled")
+
+
+def decode_vector(vec_3072, max_len=80, temp=0.7, top_p=0.9, rep_penalty=1.3):
+    if decoder_model is None:
+        return None
+    v = vec_3072.unsqueeze(0) if vec_3072.dim() == 1 else vec_3072
+    with torch.no_grad():
+        logits = decoder_model(v)
+        next_logits = logits[0, -1, :] / temp
+        probs = F.softmax(next_logits, dim=-1)
+        sp, si = torch.sort(probs, descending=True)
+        cp = sp.cumsum(0)
+        sp[cp - sp > top_p] = 0
+        sp = sp / sp.sum()
+        first = si[torch.multinomial(sp, 1)]
+
+    generated = [first.item()]
+    for _ in range(max_len - 1):
+        tok_in = torch.tensor([generated], dtype=torch.long, device=DEV)
+        with torch.no_grad():
+            logits = decoder_model(v, tok_in)
+            next_logits = logits[0, -1, :]
+            for t in set(generated[-64:]):
+                next_logits[t] /= rep_penalty
+            next_logits = next_logits / temp
+            probs = F.softmax(next_logits, dim=-1)
+            sp, si = torch.sort(probs, descending=True)
+            cp = sp.cumsum(0)
+            sp[cp - sp > top_p] = 0
+            sp = sp / sp.sum()
+            nxt = si[torch.multinomial(sp, 1)].item()
+        if dec_eot_id is not None and nxt == dec_eot_id:
+            break
+        generated.append(nxt)
+    return decoder_tk.decode(generated).strip()
+
+
+STRIP = re.compile(r'^\[conversation\] I replied \(puppet\):\s*["\']?', re.I)
+
+
+def navigate(query: str, top_k: int = 6, ew_boost: bool = True):
+    qe = torch.tensor(
+        embedder.encode([query], normalize_embeddings=True),
+        dtype=torch.float32
+    ).to(DEV)
+    with torch.no_grad():
+        rv = model(qe, mem_vecs)
+    sims = (mem_vecs @ rv.T).squeeze()
+
+    if ew_boost:
+        ew_t = torch.tensor(ew_list, dtype=torch.float32)
+        boost = 1.0 + 0.15 * (ew_t - 5.0) / 5.0
+        scored = sims * boost
+    else:
+        scored = sims
+
+    n_cand = min(top_k * 4, len(mem_texts))
+    cand_idx = scored.topk(n_cand).indices.tolist()
+
+    picked = []
+    for i in cand_idx:
+        if len(picked) >= top_k:
+            break
+        t = STRIP.sub("", mem_texts[i]).strip().strip('"').strip("'")
+        t_short = t[:200].lower()
+        too_similar = False
+        for prev_t, _ in picked:
+            overlap = len(set(t_short.split()) & set(prev_t.split())) / max(len(set(t_short.split())), 1)
+            if overlap > 0.6:
+                too_similar = True
+                break
+        if not too_similar:
+            picked.append((t_short, i))
+
+    results = []
+    for _, i in picked:
+        t = STRIP.sub("", mem_texts[i]).strip().strip('"').strip("'")
+        results.append({
+            "text": t[:400],
+            "emotional_weight": ew_list[i],
+            "salience": sal_list[i],
+            "similarity": float(sims[i]),
+            "score": float(scored[i]),
+            "idx": i
+        })
+    return results
+
+
+def raw_cosine(query: str, top_k: int = 4):
+    qe = torch.tensor(
+        embedder.encode([query], normalize_embeddings=True),
+        dtype=torch.float32
+    )
+    sims = (mem_mini @ qe.T).squeeze()
+    top_idx = sims.topk(top_k).indices.tolist()
+    results = []
+    for i in top_idx:
+        t = STRIP.sub("", mem_texts[i]).strip().strip('"').strip("'")
+        results.append({
+            "text": t[:400],
+            "emotional_weight": ew_list[i],
+            "salience": sal_list[i],
+            "similarity": float(sims[i]),
+            "idx": i,
+            "source": "cosine"
+        })
+    return results
+
+
+def blend(query: str, top_k: int = 6):
+    nav_results = navigate(query, top_k=top_k, ew_boost=True)
+    cos_results = raw_cosine(query, top_k=top_k)
+    for r in nav_results:
+        r["source"] = "navigator"
+
+    seen_idx = set()
+    merged = []
+    ni, ci = 0, 0
+    while len(merged) < top_k and (ni < len(nav_results) or ci < len(cos_results)):
+        for _ in range(2):
+            if ni < len(nav_results) and nav_results[ni]["idx"] not in seen_idx:
+                seen_idx.add(nav_results[ni]["idx"])
+                merged.append(nav_results[ni])
+            ni += 1
+        if ci < len(cos_results) and cos_results[ci]["idx"] not in seen_idx:
+            seen_idx.add(cos_results[ci]["idx"])
+            merged.append(cos_results[ci])
+        ci += 1
+    return merged[:top_k]
+
+
+class Handler(BaseHTTPRequestHandler):
+    def log_message(self, fmt, *args): pass
+
+    def _cors(self):
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
+        self.send_header("Access-Control-Allow-Headers", "Content-Type")
+
+    def do_OPTIONS(self):
+        self.send_response(200); self._cors(); self.end_headers()
+
+    def do_POST(self):
+        if self.path not in ("/navigate", "/blend", "/attention", "/decode", "/synthesize"):
+            self.send_response(404); self.end_headers(); return
+        length = int(self.headers.get("Content-Length", 0))
+        body = json.loads(self.rfile.read(length))
+        query = body.get("query", "")
+        top_k = int(body.get("top_k", 6))
+        t0 = time.time()
+
+        if self.path == "/decode":
+            if decoder_model is None:
+                result = {"error": "decoder not loaded"}
+            else:
+                vec_data = body.get("vector")
+                if vec_data:
+                    v = torch.tensor([vec_data], dtype=torch.float32).to(DEV)
+                    v = F.normalize(v, dim=-1)
+                elif query:
+                    qe = torch.tensor(
+                        embedder.encode([query], normalize_embeddings=True),
+                        dtype=torch.float32
+                    ).to(DEV)
+                    with torch.no_grad():
+                        v = model(qe, mem_vecs)
+                else:
+                    result = {"error": "provide query or vector"}
+                    v = None
+                if v is not None:
+                    text = decode_vector(v.squeeze(0),
+                                         max_len=int(body.get("max_len", 80)),
+                                         temp=float(body.get("temperature", 0.7)))
+                    result = {"text": text}
+            elapsed = time.time() - t0
+            result["elapsed"] = elapsed
+            resp = json.dumps(result).encode()
+            self.send_response(200); self._cors()
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(resp)))
+            self.end_headers(); self.wfile.write(resp)
+            print(f"[rmm] /decode {repr(query[:40])} -> {repr((result.get('text') or '')[:60])} ({elapsed:.2f}s)")
+            return
+
+        if self.path == "/synthesize":
+            mems = blend(query, top_k)
+            synth_text = None
+            if decoder_model is not None and query:
+                qe = torch.tensor(
+                    embedder.encode([query], normalize_embeddings=True),
+                    dtype=torch.float32
+                ).to(DEV)
+                with torch.no_grad():
+                    rv = model(qe, mem_vecs)
+                synth_text = decode_vector(rv.squeeze(0),
+                                           max_len=int(body.get("max_len", 80)),
+                                           temp=float(body.get("temperature", 0.7)))
+            elapsed = time.time() - t0
+            result = {"synthesized": synth_text, "memories": mems, "elapsed": elapsed}
+            resp = json.dumps(result).encode()
+            self.send_response(200); self._cors()
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(resp)))
+            self.end_headers(); self.wfile.write(resp)
+            print(f"[rmm] /synthesize {repr(query[:40])} -> synth={repr((synth_text or '')[:60])} + {len(mems)} mems ({elapsed:.2f}s)")
+            return
+
+        if self.path == "/attention":
+            qe = torch.tensor(
+                embedder.encode([query], normalize_embeddings=True),
+                dtype=torch.float32
+            ).to(DEV)
+            with torch.no_grad():
+                rv, attn_list = model(qe, mem_vecs, return_attn=True)
+            avg_attn = torch.stack([a.squeeze(0).squeeze(0) for a in attn_list]).mean(0)
+            top_attn_idx = avg_attn.topk(top_k).indices.tolist()
+            mems = []
+            for i in top_attn_idx:
+                t = STRIP.sub("", mem_texts[i]).strip().strip('"').strip("'")
+                mems.append({
+                    "text": t[:400],
+                    "emotional_weight": ew_list[i],
+                    "attention": float(avg_attn[i]),
+                    "idx": i
+                })
+            result = {"attended": mems}
+        elif self.path == "/blend":
+            mems = blend(query, top_k)
+            result = {"memories": mems}
+        else:
+            mems = navigate(query, top_k)
+            result = {"memories": mems}
+
+        elapsed = time.time() - t0
+        result["elapsed"] = elapsed
+        resp = json.dumps(result).encode()
+        self.send_response(200); self._cors()
+        self.send_header("Content-Type", "application/json")
+        self.send_header("Content-Length", str(len(resp)))
+        self.end_headers(); self.wfile.write(resp)
+        print(f"[rmm] {self.path} {repr(query[:40])} -> {len(mems)} results ({elapsed:.2f}s)")
+
+    def do_GET(self):
+        if self.path == "/health":
+            resp = b'{"status":"ok"}'
+            self.send_response(200); self._cors()
+            self.send_header("Content-Type","application/json")
+            self.send_header("Content-Length",str(len(resp)))
+            self.end_headers(); self.wfile.write(resp)
+
+
+class ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
+    daemon_threads = True
+
+if __name__ == "__main__":
+    server = ThreadedHTTPServer(("0.0.0.0", PORT), Handler)
+    print(f"[rmm] listening on http://localhost:{PORT}")
+    print(f"[rmm] endpoints: /navigate /blend /decode /synthesize /attention /health")
+    server.serve_forever()