| """ |
| 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 = {} |
|
|
| |
| 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 |
|
|
| |
| @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 |
|
|
