Upload train_navigator.py with huggingface_hub

train_navigator.py

@@ -0,0 +1,298 @@
+"""train_navigator.py — Train the RMM Navigator.
+
+Architecture:
+  Query (text) -> MiniLM embed (384-d) -> Linear -> 3072-d query vector
+  Cross-attention over entity's spine memory vectors (3072-d each)
+  -> Synthesized response vector (3072-d)
+  -> Cosine loss vs actual reply vector from spine
+
+The navigator learns the topology of an entity's embedding space —
+which memories are connected, which regions respond to which queries,
+how emotional weight shapes retrieval. This is learned navigation,
+not cosine similarity.
+
+Run: modal run train_navigator.py
+Pull: modal volume get rmm-vol /memory-nav/ ./memory-nav-out/
+
+Requires:
+  - spine.json: {"memories": [{"text": "...", "vector": [...3072...], "emotional_weight": 8, "salience": 0.5}, ...]}
+  - dialogue.txt: alternating "Speaker A: ...\nSpeaker B: ..." blocks
+  - (optional) discord.json: array of {author: {name: "..."}, content: "..."} messages
+"""
+import modal, json
+from pathlib import Path
+
+app = modal.App("rmm-navigator")
+image = (modal.Image.debian_slim(python_version="3.11")
+         .pip_install("torch==2.6.0", "numpy", "sentence-transformers"))
+vol = modal.Volume.from_name("rmm-vol", create_if_missing=True)
+
+# ── Point these at your entity's data ──
+SPINE_FILE    = Path("spine.json")
+DIALOGUE_FILE = Path("dialogue.txt")
+DISCORD_FILE  = Path("discord.json")
+
+SPINE_DIM  = 3072   # embedding dim (Gemini, etc.)
+QUERY_DIM  = 384    # MiniLM dim
+N_HEADS    = 8
+N_LAYERS   = 3
+D_MODEL    = 512
+DROPOUT    = 0.1
+
+
+@app.function(image=image, gpu="A10G", timeout=3600, volumes={"/vol": vol})
+def train(spine_json: str, dialogue_text: str, discord_json: str = "",
+          speaker_a: str = "Laura", speaker_b: str = "Micah", smoke: bool = False):
+    import os, math, time, json, re
+    import numpy as np
+    import torch
+    import torch.nn as nn
+    import torch.nn.functional as F
+    from sentence_transformers import SentenceTransformer
+
+    DEV = "cuda"
+    print(f"[nav] gpu={torch.cuda.get_device_name(0)}")
+
+    spine_data = json.loads(spine_json)
+    mems = spine_data["memories"]
+
+    mem_vecs = torch.tensor(
+        [m["vector"] for m in mems], dtype=torch.float32
+    ).to(DEV)
+    mem_vecs = F.normalize(mem_vecs, dim=-1)
+    N_MEM = mem_vecs.shape[0]
+    print(f"[nav] {N_MEM} memory vectors loaded, dim={mem_vecs.shape[1]}")
+
+    STRIP = re.compile(r'^\[conversation\] I replied \(puppet\):\s*["\']?', re.I)
+    SURR = re.compile(r'[\ud800-\udfff]')
+    mem_texts = []
+    for m in mems:
+        raw = SURR.sub('', str(m.get("text") or ""))
+        t = STRIP.sub("", raw).strip().strip('"').strip("'")
+        mem_texts.append(t[:500] if t else "...")
+
+    saliences = torch.tensor(
+        [m.get("salience", 0.5) for m in mems], dtype=torch.float32
+    ).to(DEV)
+
+    print("[nav] building training pairs...")
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+
+    query_msgs, reply_msgs = [], []
+
+    # Source 1: dialogue file (SpeakerA: ...\nSpeakerB: ... blocks)
+    blocks = [b.strip() for b in dialogue_text.split("\n\n")
+              if speaker_a + ":" in b and speaker_b + ":" in b]
+    n_smoke = 50 if smoke else len(blocks)
+    for b in blocks[:n_smoke]:
+        parts = b.split(speaker_b + ":", 1)
+        a_part = parts[0].replace(speaker_a + ":", "").strip()
+        b_part = parts[1].strip() if len(parts) > 1 else ""
+        if len(a_part) >= 5 and len(b_part) >= 10:
+            query_msgs.append(a_part)
+            reply_msgs.append(b_part)
+    print(f"[nav] dialogue file: {len(query_msgs)} pairs")
+
+    # Source 2: Discord history (consecutive A->B messages)
+    if discord_json and not smoke:
+        disc_msgs = json.loads(discord_json)
+        disc_count = 0
+        for i in range(len(disc_msgs) - 1):
+            cur = disc_msgs[i]
+            nxt = disc_msgs[i + 1]
+            cur_a = (cur.get("author", {}).get("name", "") or cur.get("author", {}).get("username", "") or "").lower()
+            nxt_a = (nxt.get("author", {}).get("name", "") or nxt.get("author", {}).get("username", "") or "").lower()
+            cur_c = SURR.sub('', cur.get("content", "").strip())
+            nxt_c = SURR.sub('', nxt.get("content", "").strip())
+            is_a = speaker_a.lower() in cur_a
+            is_b = speaker_b.lower() in nxt_a
+            if is_a and is_b and len(cur_c) >= 3 and len(nxt_c) >= 10:
+                query_msgs.append(cur_c[:500])
+                reply_msgs.append(nxt_c[:500])
+                disc_count += 1
+        print(f"[nav] discord pairs: {disc_count}")
+
+    # Sanitize
+    clean_q, clean_r = [], []
+    for q, r in zip(query_msgs, reply_msgs):
+        qs, rs = str(q).strip(), str(r).strip()
+        if len(qs) >= 3 and len(rs) >= 5:
+            clean_q.append(qs)
+            clean_r.append(rs)
+    query_msgs, reply_msgs = clean_q, clean_r
+
+    print(f"[nav] total pairs: {len(query_msgs)}")
+    print(f"[nav] embedding queries...")
+    query_embs = embedder.encode(query_msgs, normalize_embeddings=True,
+                                  show_progress_bar=False, batch_size=64)
+
+    print(f"[nav] embedding replies...")
+    reply_embs = embedder.encode(reply_msgs, normalize_embeddings=True,
+                                  show_progress_bar=False, batch_size=64)
+
+    reply_tensor = torch.tensor(reply_embs, dtype=torch.float32)
+
+    print("[nav] embedding memories in MiniLM space for matching...")
+    BATCH = 256
+    mem_mini_embs = []
+    for start in range(0, N_MEM, BATCH):
+        chunk = mem_texts[start:start + BATCH]
+        e = embedder.encode(chunk, normalize_embeddings=True, show_progress_bar=False)
+        mem_mini_embs.append(e)
+    mem_mini = torch.tensor(np.vstack(mem_mini_embs), dtype=torch.float32)
+
+    sims = reply_tensor @ mem_mini.T
+    top5_vals, top5_idx = sims.topk(5, dim=-1)
+    sal_cpu = saliences.cpu()
+    best_indices = []
+    for i in range(len(reply_tensor)):
+        candidates = top5_idx[i]
+        cand_sals = sal_cpu[candidates]
+        best_j = cand_sals.argmax().item()
+        best_indices.append(candidates[best_j].item())
+    best_mem_idx = torch.tensor(best_indices, dtype=torch.long)
+    target_vecs = mem_vecs[best_mem_idx]
+
+    ew_raw = torch.tensor([mems[i].get("emotional_weight", 5) for i in best_indices],
+                          dtype=torch.float32)
+    pair_weights = 1.0 + 0.3 * (ew_raw - 5.0) / 5.0
+    pair_weights = pair_weights / pair_weights.mean()
+
+    query_tensor = torch.tensor(query_embs, dtype=torch.float32)
+    print(f"[nav] {len(query_tensor)} training pairs ready")
+
+    # ── Model ──
+    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=DROPOUT, 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):
+            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
+            for layer in self.layers:
+                x = layer(x, m)
+            x = self.norm(x).squeeze(1)
+            out = self.out_proj(x)
+            return F.normalize(out, dim=-1)
+
+    model = MemoryNavigator().to(DEV)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[nav] model {n_params/1e6:.1f}M params")
+
+    # ── Train ──
+    ITERS = 200 if smoke else 7500
+    BS = 32
+    N_NEG = 7
+    MARGIN = 0.2
+    opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
+    warmup_steps = 200 if not smoke else 20
+    def lr_lambda(step):
+        if step < warmup_steps:
+            return step / warmup_steps
+        progress = (step - warmup_steps) / max(1, ITERS - warmup_steps)
+        return 0.5 * (1 + math.cos(math.pi * progress))
+    sch = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
+
+    M = len(query_tensor)
+    t0 = time.time()
+    best_loss = float('inf')
+    best_state = None
+
+    for step in range(ITERS):
+        idx = torch.randint(0, M, (BS,))
+        q_batch = query_tensor[idx].to(DEV)
+        t_batch = target_vecs[idx].to(DEV)
+        t_idx = best_mem_idx[idx]
+
+        pred = model(q_batch, mem_vecs)
+
+        pos_sim = (pred * t_batch).sum(dim=-1)
+
+        neg_sims_list = []
+        for b in range(BS):
+            all_sims = (mem_vecs @ pred[b]).squeeze()
+            all_sims[t_idx[b]] = -1.0
+            hard_neg_idx = all_sims.topk(N_NEG).indices
+            neg_sims_list.append(all_sims[hard_neg_idx].mean())
+        neg_sim = torch.stack(neg_sims_list)
+
+        w = pair_weights[idx].to(DEV)
+        loss_pos = ((1.0 - pos_sim) * w).mean()
+        loss_neg = (F.relu(neg_sim - pos_sim + MARGIN) * w).mean()
+        loss = loss_pos + 0.3 * loss_neg
+
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+        sch.step()
+
+        if step % (20 if smoke else 250) == 0:
+            lv = loss.item()
+            lp = loss_pos.item()
+            ln = loss_neg.item()
+            mark = " <-" if lv < best_loss else ""
+            print(f"  [nav] step {step:4d} loss={lv:.4f} (pos={lp:.4f} neg={ln:.4f}) ({time.time()-t0:.0f}s){mark}")
+            if lv < best_loss:
+                best_loss = lv
+                best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+
+    if best_state:
+        model.load_state_dict(best_state)
+
+    os.makedirs("/vol/memory-nav", exist_ok=True)
+    torch.save({k: v.cpu() for k, v in model.state_dict().items()},
+               "/vol/memory-nav/navigator.pt")
+    np.save("/vol/memory-nav/mem_vecs.npy", mem_vecs.cpu().numpy())
+    import pickle
+    with open("/vol/memory-nav/mem_texts.pkl", "wb") as f:
+        pickle.dump(mem_texts, f)
+    vol.commit()
+    print(f"[nav] DONE best_loss={best_loss:.4f} saved to /vol/memory-nav/")
+
+    model.eval()
+    test_queries = ["hello", "I love you", "I miss her", "tell me a story"]
+    for q in test_queries:
+        qe = torch.tensor(
+            embedder.encode([q], normalize_embeddings=True), dtype=torch.float32
+        ).to(DEV)
+        with torch.no_grad():
+            rv = model(qe, mem_vecs)
+        sims = (mem_vecs @ rv.T).squeeze()
+        top3 = sims.topk(3).indices.tolist()
+        print(f"\nQuery: {q!r}")
+        for i in top3:
+            print(f"  [{i}] ew={mems[i].get('emotional_weight',0)} {mem_texts[i][:100]}")
+
+    return {"best_loss": best_loss, "params_m": n_params/1e6}
+
+
+@app.local_entrypoint()
+def main(smoke: bool = False):
+    spine_json = SPINE_FILE.read_text(encoding="utf-8", errors="ignore")
+    dialogue   = DIALOGUE_FILE.read_text(encoding="utf-8", errors="ignore")
+    discord = ""
+    if DISCORD_FILE.exists() and not smoke:
+        discord = DISCORD_FILE.read_text(encoding="utf-8", errors="ignore")
+    n_pairs = len([b for b in dialogue.split(chr(10)*2) if 'Laura:' in b])
+    print(f"[local] spine={len(spine_json)//1024}KB  dialogue={n_pairs}  discord={len(discord)//1024}KB  smoke={smoke}")
+    r = train.remote(spine_json, dialogue, discord, smoke=smoke)
+    print(f"[local] done  loss={r['best_loss']:.4f}  params={r['params_m']:.1f}M")