"""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")