File size: 4,093 Bytes
933929e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """
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
|