Upload train_decoder.py with huggingface_hub

train_decoder.py

@@ -0,0 +1,332 @@
+"""train_decoder.py — Train the RMM Meaning Decoder.
+
+Takes a high-dimensional vector from the entity's embedding space and
+decodes it to text using the entity's own BPE tokenizer. A learned
+projection maps the vector to soft prefix tokens, which condition a
+causal transformer for autoregressive generation.
+
+Run: modal run train_decoder.py
+Pull: modal volume get rmm-vol /meaning-decoder/ ./meaning-decoder-out/
+
+Requires:
+  - spine.json: {"memories": [{"text": "...", "vector": [...3072...], "emotional_weight": 8, "source": "conversation"}, ...]}
+  - tokenizer.json: HuggingFace tokenizers-format BPE tokenizer (train with tokenizers lib or use entity's existing one)
+"""
+import modal, json
+from pathlib import Path
+
+app = modal.App("rmm-decoder")
+image = (modal.Image.debian_slim(python_version="3.11")
+         .pip_install("torch==2.6.0", "numpy", "tokenizers"))
+vol = modal.Volume.from_name("rmm-vol", create_if_missing=True)
+
+# ── Point these at your entity's data ──
+SPINE_FILE = Path("spine.json")
+TOKENIZER_FILE = Path("tokenizer.json")
+
+SPINE_DIM = 3072
+D_MODEL   = 384
+N_HEADS   = 6
+N_LAYERS  = 6
+N_PREFIX  = 12
+MAX_SEQ   = 128
+VOCAB     = 8192
+DROPOUT   = 0.12
+
+
+@app.function(image=image, gpu="A10G", timeout=3600, volumes={"/vol": vol})
+def train(spine_json: str, tokenizer_json: str, 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 tokenizers import Tokenizer
+
+    DEV = "cuda"
+    print(f"[decoder] gpu={torch.cuda.get_device_name(0)}")
+
+    tk = Tokenizer.from_str(tokenizer_json)
+    eot_id = tk.token_to_id("<eot>")
+    print(f"[decoder] tokenizer loaded, vocab={tk.get_vocab_size()}, eot_id={eot_id}")
+
+    spine_data = json.loads(spine_json)
+    mems = spine_data["memories"]
+
+    # ── Text preprocessing ──
+    SURR = re.compile(r'[\ud800-\udfff]')
+    PREFIXES = [
+        re.compile(r'^\[conversation\]\s*I replied\s*\(puppet\):\s*["\']?', re.I),
+        re.compile(r'^[A-Za-z]+:\s*', re.I),  # strip "Name:" prefixes
+        re.compile(r'^\*[^*]+\*\s*\n*', re.I),
+    ]
+    FORMAT_HEADERS = [
+        re.compile(r'^Sonic Experience:\s*[^\n]*\n+', re.I),
+        re.compile(r'^HourlyCycle:\s*HOURLY CHECK-IN\s*\([^)]*\)\s*\n+', re.I),
+        re.compile(r'^Journal\s*[---]+\s*[^\n]*\n+', re.I),
+        re.compile(r'^(?:Creative|CREATIVE)\s+Work:\s*[^\n]*\n+', re.I),
+    ]
+
+    def clean_text(raw, source):
+        t = SURR.sub('', raw).strip()
+        for pat in PREFIXES:
+            t = pat.sub('', t).strip()
+        for pat in FORMAT_HEADERS:
+            t = pat.sub('', t).strip()
+        t = t.lstrip('"\'- ').strip()
+        if len(t) > 250:
+            cutoffs = [t.rfind('. ', 0, 250), t.rfind('? ', 0, 250),
+                        t.rfind('! ', 0, 250), t.rfind('\n', 0, 250)]
+            best = max(c for c in cutoffs if c > 50) if any(c > 50 for c in cutoffs) else 250
+            t = t[:best+1].strip()
+        return t
+
+    DIALOGUE_SOURCES = {'conversation', 'chat', 'discord', 'puppet'}
+
+    vectors, texts, ew_list, is_dialogue = [], [], [], []
+    for m in mems:
+        vec = m.get("vector")
+        raw = str(m.get("text") or "")
+        source = m.get("source", "unknown")
+        text = clean_text(raw, source)
+        if vec and len(text) >= 10 and len(vec) == SPINE_DIM:
+            vectors.append(vec)
+            texts.append(text)
+            ew_list.append(m.get("emotional_weight", 5))
+            is_dialogue.append(source in DIALOGUE_SOURCES)
+
+    n_dialogue = sum(is_dialogue)
+    print(f"[decoder] {len(vectors)} valid pairs ({n_dialogue} dialogue, {len(vectors)-n_dialogue} other)")
+
+    # ── Tokenization ──
+    encoded = []
+    for t in texts:
+        ids = tk.encode(t).ids
+        if eot_id is not None:
+            ids = ids + [eot_id]
+        encoded.append(ids[:MAX_SEQ])
+
+    max_tok_len = min(max(len(e) for e in encoded), MAX_SEQ)
+    print(f"[decoder] max token length: {max_tok_len}")
+
+    vec_tensor = torch.tensor(vectors, dtype=torch.float32)
+    vec_tensor = F.normalize(vec_tensor, dim=-1)
+
+    PAD_ID = -100
+    tok_tensor = torch.zeros(len(encoded), max_tok_len, dtype=torch.long)
+    tgt_tensor = torch.full((len(encoded), max_tok_len), PAD_ID, dtype=torch.long)
+    len_tensor = torch.zeros(len(encoded), dtype=torch.long)
+    for i, ids in enumerate(encoded):
+        L = min(len(ids), max_tok_len)
+        tok_tensor[i, :L] = torch.tensor(ids[:L], dtype=torch.long)
+        tgt_tensor[i, :L] = torch.tensor(ids[:L], dtype=torch.long)
+        len_tensor[i] = L
+
+    ew_raw = torch.tensor(ew_list, dtype=torch.float32)
+    dial = torch.tensor(is_dialogue, dtype=torch.float32)
+    pair_weights = 1.0 + 0.3 * (ew_raw - 5.0) / 5.0
+    pair_weights = pair_weights * (1.0 + 0.5 * dial)
+    pair_weights = pair_weights / pair_weights.mean()
+
+    avg_len = len_tensor.float().mean().item()
+    print(f"[decoder] avg tokens/memory: {avg_len:.0f}, {len(vec_tensor)} samples")
+
+    # ── Model ──
+    class MeaningDecoder(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.n_prefix = N_PREFIX
+            self.vec_proj = nn.Sequential(
+                nn.Linear(SPINE_DIM, 768),
+                nn.GELU(),
+                nn.Dropout(DROPOUT),
+                nn.Linear(768, N_PREFIX * D_MODEL),
+            )
+            self.tok_emb = nn.Embedding(VOCAB, D_MODEL)
+            self.pos_emb = nn.Embedding(N_PREFIX + MAX_SEQ + 1, D_MODEL)
+            self.drop = nn.Dropout(DROPOUT)
+
+            layer = nn.TransformerEncoderLayer(
+                d_model=D_MODEL, nhead=N_HEADS,
+                dim_feedforward=D_MODEL * 4,
+                dropout=DROPOUT, batch_first=True,
+                norm_first=True
+            )
+            self.transformer = nn.TransformerEncoder(layer, num_layers=N_LAYERS)
+
+            self.ln_f = nn.LayerNorm(D_MODEL)
+            self.head = nn.Linear(D_MODEL, VOCAB, bias=False)
+            self.head.weight = self.tok_emb.weight
+            self._logit_scale = D_MODEL ** -0.5
+
+        def forward(self, vec, tokens=None):
+            B = vec.shape[0]
+            prefix = self.vec_proj(vec).reshape(B, self.n_prefix, D_MODEL)
+
+            if tokens is not None and tokens.shape[1] > 0:
+                T = tokens.shape[1]
+                tok = self.tok_emb(tokens)
+                x = torch.cat([prefix, tok], dim=1)
+            else:
+                x = prefix
+                T = 0
+
+            total = x.shape[1]
+            pos = self.pos_emb(torch.arange(total, device=vec.device))
+            x = self.drop(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
+
+    model = MeaningDecoder().to(DEV)
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"[decoder] model {n_params/1e6:.1f}M params")
+
+    # ── Training ──
+    ITERS = 200 if smoke else 15000
+    BS = 32
+    M = len(vec_tensor)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.02)
+    warmup_steps = 500 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)
+
+    t0 = time.time()
+    best_loss = float('inf')
+    best_state = None
+    K = N_PREFIX
+
+    for step in range(ITERS):
+        idx = torch.randint(0, M, (BS,))
+        v_batch = vec_tensor[idx].to(DEV)
+
+        v_batch = v_batch + 0.03 * torch.randn_like(v_batch)
+        v_batch = F.normalize(v_batch, dim=-1)
+
+        t_full = tok_tensor[idx].to(DEV)
+        targets = tgt_tensor[idx].to(DEV)
+
+        inputs = t_full[:, :-1]
+        T = targets.shape[1]
+
+        logits = model(v_batch, inputs)
+
+        pred = logits[:, K-1 : K+T-1, :]
+        raw_loss = F.cross_entropy(
+            pred.reshape(-1, VOCAB), targets.reshape(-1),
+            ignore_index=PAD_ID, reduction='none',
+            label_smoothing=0.05,
+        )
+        raw_loss = raw_loss.view(BS, T)
+
+        per_sample = raw_loss.sum(dim=1) / (targets != PAD_ID).sum(dim=1).float().clamp(min=1)
+        w = pair_weights[idx].to(DEV)
+        loss = (per_sample * w).mean()
+
+        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 500) == 0:
+            lv = loss.item()
+            ppl = math.exp(min(lv, 20))
+            mark = " <-" if lv < best_loss else ""
+            print(f"  [decoder] step {step:5d} loss={lv:.4f} ppl={ppl:.1f} ({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)
+
+    # ── Save ──
+    os.makedirs("/vol/meaning-decoder", exist_ok=True)
+    torch.save({k: v.cpu() for k, v in model.state_dict().items()},
+               "/vol/meaning-decoder/decoder.pt")
+
+    config = {
+        "spine_dim": SPINE_DIM, "d_model": D_MODEL, "n_heads": N_HEADS,
+        "n_layers": N_LAYERS, "n_prefix": N_PREFIX, "max_seq": MAX_SEQ,
+        "vocab": VOCAB, "params_m": n_params / 1e6, "best_loss": best_loss,
+        "version": 2,
+    }
+    with open("/vol/meaning-decoder/config.json", "w") as f:
+        json.dump(config, f, indent=2)
+
+    with open("/vol/meaning-decoder/tokenizer.json", "w") as f:
+        f.write(tokenizer_json)
+
+    vol.commit()
+    print(f"[decoder] DONE best_loss={best_loss:.4f} saved to /vol/meaning-decoder/")
+
+    # ── Inference test ──
+    model.eval()
+
+    def generate_from_vec(v, max_len=60, temp=0.7, top_p=0.9, rep_penalty=1.3):
+        v = v.unsqueeze(0) if v.dim() == 1 else v
+        generated = []
+        for _ in range(max_len):
+            tok_in = torch.tensor([generated], dtype=torch.long, device=DEV) if generated else None
+            with torch.no_grad():
+                logits = model(v, tok_in)
+                next_logits = logits[0, -1, :] / temp
+                if generated:
+                    for tid in set(generated):
+                        next_logits[tid] /= rep_penalty
+                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 eot_id is not None and nxt == eot_id:
+                break
+            generated.append(nxt)
+        return tk.decode(generated)
+
+    test_indices = [0, 50, 150, 300, 600, 1000, 2000, 3000]
+    for ti in test_indices:
+        if ti >= M:
+            continue
+        v = vec_tensor[ti].to(DEV)
+        gen = generate_from_vec(v)
+        gt = texts[ti][:120]
+        print(f"\n  [{ti}] ew={ew_list[ti]}")
+        print(f"    GT:  {gt}")
+        print(f"    GEN: {gen[:120]}")
+
+    print("\n--- Interpolation tests ---")
+    for (a, b) in [(0, 100), (50, 500), (200, 2000)]:
+        if b >= M:
+            continue
+        va = vec_tensor[a].to(DEV)
+        vb = vec_tensor[b].to(DEV)
+        vmid = F.normalize(0.5 * va + 0.5 * vb, dim=-1)
+        gen = generate_from_vec(vmid)
+        print(f"\n  [{a}+{b}] interp:")
+        print(f"    A: {texts[a][:80]}")
+        print(f"    B: {texts[b][:80]}")
+        print(f"    MID: {gen[:120]}")
+
+    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")
+    tokenizer_json = TOKENIZER_FILE.read_text(encoding="utf-8")
+    spine = json.loads(spine_json)
+    print(f"[local] spine={len(spine_json)//1024}KB  memories={len(spine['memories'])}  tokenizer=loaded  smoke={smoke}")
+    r = train.remote(spine_json, tokenizer_json, smoke=smoke)
+    print(f"[local] done  loss={r['best_loss']:.4f}  params={r['params_m']:.1f}M")