Upload memstore.py with huggingface_hub

memstore.py

@@ -0,0 +1,112 @@
+"""
+clankerDiffusion — learned secondary memory (side store).
+
+The model controls a long-term store through special tokens it emits while
+generating:
+
+    <mem_write>KEY<mem_kv>BLOB</mem_kv>      -> persist BLOB under KEY
+    <mem_read>KEY</mem_read>                -> recall BLOB for KEY
+    <mem_evict>KEY</mem_evict>              -> delete KEY
+
+BLOBs are packed into the token stream as <mem_kv>...</mem_kv> so the model can
+read what it wrote. At inference, `MemoryStore` intercepts these tokens, performs
+the read/write against an external store (dict / JSONL / vector index), and
+replaces the in-stream content with the resolved value (or an <unk> if missing).
+
+This is "knowledge brought in and out of memory": the weights never hold the
+facts; only the keys/pointers do. Train 8k, but the store is unbounded, so the
+effective secondary memory is far larger than the 128k normal window.
+
+Training note: the data generator emits syntactically-valid mem_write / mem_read
+/ mem_evict spans using this same module, so the model learns the format and
+when to invoke it.
+"""
+import re
+import json
+import os
+
+
+class MemoryStore:
+    def __init__(self, path=None):
+        self.path = path
+        self.store = {}
+        if path and os.path.exists(path):
+            try:
+                with open(path, "r", encoding="utf-8") as f:
+                    self.store = json.load(f)
+            except Exception:
+                self.store = {}
+
+    # ---- token-stream helpers (used by data gen + inference) ----
+    def write(self, key, blob):
+        self.store[key] = blob
+        self._maybe_persist()
+        return blob
+
+    def read(self, key):
+        return self.store.get(key)
+
+    def evict(self, key):
+        return self.store.pop(key, None)
+
+    def _maybe_persist(self):
+        if self.path:
+            try:
+                with open(self.path, "w", encoding="utf-8") as f:
+                    json.dump(self.store, f)
+            except Exception:
+                pass
+
+    # ---- stream rewriting for inference ----
+    @staticmethod
+    def _split_segments(text):
+        """Yield (is_mem, content) chunks where is_mem=True spans a full
+        <mem_*>...</mem_*> block we can act on."""
+        pat = re.compile(
+            r"<mem_write>(.*?)<mem_kv>(.*?)</mem_kv>|"
+            r"<mem_read>(.*?)</mem_read>|"
+            r"<mem_evict>(.*?)</mem_evict>",
+            re.DOTALL)
+        pos = 0
+        for m in pat.finditer(text):
+            if m.start() > pos:
+                yield (False, text[pos:m.start()])
+            if m.group(1) is not None:
+                key, blob = m.group(1), m.group(2)
+                yield (True, ("write", key, blob))
+            elif m.group(3) is not None:
+                yield (True, ("read", m.group(3)))
+            elif m.group(4) is not None:
+                yield (True, ("evict", m.group(4)))
+            pos = m.end()
+        if pos < len(text):
+            yield (False, text[pos:])
+
+    def resolve(self, text):
+        """Rewrite a generated stream: execute memory ops, replacing the op
+        with its resolved <mem_kv> value (or a miss marker). Returns the new
+        text and a list of (op, key) that were performed (for logging)."""
+        out = []
+        ops = []
+        for is_mem, chunk in self._split_segments(text):
+            if not is_mem:
+                out.append(chunk)
+                continue
+            op = chunk[0]
+            if op == "write":
+                _, key, blob = chunk
+                self.write(key, blob)
+                ops.append(("write", key))
+                out.append(f"<mem_kv>{blob}</mem_kv>")
+            elif op == "read":
+                key = chunk[1]
+                blob = self.read(key)
+                ops.append(("read", key))
+                out.append(f"<mem_kv>{blob}</mem_kv>" if blob is not None
+                           else "<mem_kv><unk></mem_kv>")
+            elif op == "evict":
+                key = chunk[1]
+                self.evict(key)
+                ops.append(("evict", key))
+                out.append("")
+        return "".join(out), ops