cortexm/vsa/palace.py
Rama main · réplica de lectura del repositorio original.
"""The Memory Palace — VSA hologram store with codec-quantized vectors.
One hologram per fact: role-bound subject/relation/value fillers plus a
λ-weighted lexical superposition (see vsa.ops). Storage is codec-
quantized (INT8 / Binary / RaBitQ / PQ), persisted as BLOBs in SQLite
and mirrored in RAM as packed numpy matrices for microsecond scoring.
A page-clustered tree index provides O(log N) retrieval once the
collection crosses ``index_threshold``; below it, flat scan wins.
Also hosts the self-healing machinery: per-record hash checks, TMR
majority vote (binary codec), and re-encoding from the symbolic Trace
when corruption exceeds the correction radius.
"""
from __future__ import annotations
import base64
import numpy as np
from cortexm.config import Config
from cortexm.errors import CodecError, StoreError
from cortexm.text.embedder import HashingEmbedder
from cortexm.trace.fact import Fact
from cortexm.trace.store import TraceStore
from cortexm.util import iso
import datetime as _dt
from cortexm.vsa.codecs import make_codec, PQCodec, Int8Codec
from cortexm.vsa.index import TreeIndex
from cortexm.vsa.ops import VSA
VEC_TABLE = """
CREATE TABLE IF NOT EXISTS vectors (
fact_id TEXT PRIMARY KEY, record BLOB NOT NULL,
vec_hash TEXT DEFAULT '', ts TEXT DEFAULT ''
)
"""
class MemoryPalace:
def __init__(self, config: Config, store: TraceStore, embedder: HashingEmbedder | None = None) -> None:
self.cfg = config
self.store = store
self.dims = config.dims
self.hasher = store.hasher
self.codec = make_codec(config.codec, config.dims, config.seed,
tmr=config.tmr)
self.vsa = VSA(config.dims, config.vsa_mode, config.seed,
config.lexical_lambda)
# Allow injecting a shared embedder for FIX 2 (persistent worker embedder)
self.embedder = embedder or HashingEmbedder(config.dims, config.seed)
store.conn.execute(VEC_TABLE)
store.conn.commit()
self._ids: list[str] = []
self._id2row: dict[str, int] = {}
self._n = 0
self._cap = 0
self._packed: np.ndarray | None = None
self._aux: np.ndarray | None = None
self._index: TreeIndex | None = None
self._index_valid = False
self._pq_buffer: list[tuple[str, np.ndarray]] = []
self.searches_flat = 0
self.searches_indexed = 0
self._load_pq_codebooks()
self._load()
# ------------------------------------------------------------- loading
def _load(self) -> None:
rows = self.store.conn.execute(
"SELECT fact_id, record FROM vectors").fetchall()
if not rows:
return
recs = [r["record"] for r in rows]
self._ids = [r["fact_id"] for r in rows]
self._id2row = {fid: i for i, fid in enumerate(self._ids)}
packed_rows = [self.codec.from_bytes(b) for b in recs]
if isinstance(self.codec, Int8Codec):
q = np.stack([p[0] for p in packed_rows])
self._aux = np.stack([np.float16(p[1]) for p in packed_rows])
self._packed = q
self._cap = self._n = len(q)
else:
self._packed = np.stack(packed_rows)
self._cap = self._n = len(self._packed)
self._index_valid = False
def _load_pq_codebooks(self) -> None:
if not isinstance(self.codec, PQCodec):
return
blob = self.store.kv_get("PQ_CODEBOOKS")
if blob:
arr = np.frombuffer(base64.b64decode(blob), dtype=np.float32)
m = self.codec.m
ks = arr.size // (m * self.codec.sub)
self.codec.set_codebooks(arr.reshape(m, ks, self.codec.sub))
def _save_pq_codebooks(self) -> None:
if isinstance(self.codec, PQCodec) and self.codec.trained:
self.store.kv_set(
"PQ_CODEBOOKS",
base64.b64encode(self.codec.codebooks.astype(np.float32).tobytes()).decode())
# ------------------------------------------------------------- writing
def encode_fact(self, fact: Fact) -> np.ndarray:
return self.vsa.encode_fact(
self.embedder.embed(fact.subject),
self.embedder.embed(fact.relation),
self.embedder.embed(fact.value))
def _grow(self, need: int) -> None:
w = self._row_width()
new_cap = max(1024, self._cap * 2, need)
if isinstance(self.codec, Int8Codec):
packed = np.zeros((new_cap, w), dtype=np.int8)
aux = np.zeros(new_cap, dtype=np.float16)
if self._packed is not None and self._n:
packed[: self._n] = self._packed[: self._n]
aux[: self._n] = self._aux[: self._n]
self._packed, self._aux = packed, aux
else:
packed = np.zeros((new_cap, w), dtype=np.uint8)
if self._packed is not None and self._n:
packed[: self._n] = self._packed[: self._n]
self._packed = packed
self._cap = new_cap
def _row_width(self) -> int:
if isinstance(self.codec, Int8Codec):
return self.dims # scale lives in the aux array
if isinstance(self.codec, PQCodec):
return self.codec.m # 8 code bytes
return self.codec.bytes_per_vector # binary / rabitq packed words
def add(self, fact_id: str, vec: np.ndarray) -> None:
if isinstance(self.codec, PQCodec) and not self.codec.trained:
self._pq_buffer.append((fact_id, vec))
if len(self._pq_buffer) >= self.codec.train_threshold:
self._flush_pq()
return
self._append(fact_id, self._encode_row(vec), vec)
def _encode_row(self, vec: np.ndarray):
if isinstance(self.codec, Int8Codec):
q = self.codec.encode_packed(vec)
return q, self.codec.encode_scale(vec)
return self.codec.encode_packed(vec), None
def _append(self, fact_id: str, row, vec) -> None:
packed_row, scale = row
blob = (self.codec.to_bytes(packed_row, scale)
if isinstance(self.codec, Int8Codec)
else self.codec.to_bytes(packed_row))
self.store.conn.execute(
"INSERT OR REPLACE INTO vectors(fact_id, record, vec_hash, ts) VALUES(?,?,?,?)",
(fact_id, blob, self.hasher.hash_bytes(blob), iso(_dt.datetime.now(_dt.timezone.utc))))
if not self.store.batching:
self.store.conn.commit()
if self._n >= self._cap:
self._grow(self._n + 1)
if isinstance(self.codec, Int8Codec):
self._packed[self._n] = packed_row
self._aux[self._n] = scale
else:
self._packed[self._n] = packed_row
if fact_id in self._id2row:
old = self._id2row[fact_id]
# overwrite in place; keep id mapping
self._ids[old] = fact_id
else:
self._id2row[fact_id] = self._n
self._ids.append(fact_id)
self._n += 1
self._index_valid = False
def add_many(self, pairs: list[tuple[str, np.ndarray]]) -> None:
for fid, v in pairs:
self.add(fid, v)
def remove_ids(self, fact_ids: list[str]) -> int:
"""Hard-remove vectors (GDPR erasure). Rebuilds packed state from
the surviving rows. Returns the number of vectors removed."""
if not fact_ids:
return 0
gone = set(fact_ids)
cur = self.store.conn.execute(
"SELECT fact_id FROM vectors").fetchall()
existing = {r["fact_id"] for r in cur}
removed = len(gone & existing)
if removed:
qmarks = ",".join("?" * len(gone))
self.store.conn.execute(
f"DELETE FROM vectors WHERE fact_id IN ({qmarks})",
list(gone))
self.store.conn.commit()
# rebuild in-memory state from disk (source of truth)
self._ids = []
self._id2row = {}
self._n = 0
self._cap = 0
self._packed = None
self._aux = None
self._index = None
self._index_valid = False
self._load()
return removed
def _flush_pq(self) -> None:
if not self._pq_buffer or not isinstance(self.codec, PQCodec):
return
if not self.codec.trained:
vecs = np.stack([v for _, v in self._pq_buffer])
self.codec.train(vecs)
self._save_pq_codebooks()
for fid, v in self._pq_buffer:
self._append(fid, self._encode_row(v), v)
self._pq_buffer.clear()
def size(self) -> int:
return self._n + len(self._pq_buffer)
def has(self, fact_id: str) -> bool:
return fact_id in self._id2row or any(f == fact_id for f, _ in self._pq_buffer)
def close(self) -> None:
self._flush_pq()
if not self.store.batching:
# Idempotent teardown: if the underlying SQLite connection
# was already closed (e.g. Memory.close() called twice or
# store.dispose() ran before palace.close()), the commit()
# would raise sqlite3.ProgrammingError. Catch + ignore so
# Memory.close() is safe to call multiple times — a
# standard Python teardown idiom.
try:
self.store.conn.commit()
except Exception:
pass
# -------------------------------------------------------------- search
def _rows_getter(self, rows: np.ndarray):
if self._packed is None or self._n == 0:
return np.zeros((0, 1), dtype=np.uint8), None
packed = self._packed[rows]
aux = self._aux[rows] if self._aux is not None else None
return packed, aux
def ensure_index(self, force: bool = False) -> None:
if self._n < self.cfg.index_threshold:
return
if self._index is not None and self._index_valid and not force:
return
self._flush_pq()
index = TreeIndex(self.codec, self._rows_getter, self._n,
branch=self.cfg.index_branch,
leaf=self.cfg.index_leaf_size,
seed=self.cfg.seed)
index.build()
self._index = index
self._index_valid = True
def search(self, q: np.ndarray, k: int = 10,
candidate_ids: set[str] | None = None) -> list[tuple[str, float]]:
self._flush_pq()
if self._n == 0:
return []
use_index = (self._n >= self.cfg.index_threshold and self._index_valid
and (candidate_ids is None or len(candidate_ids) >= self._n * 0.5))
if use_index:
self.searches_indexed += 1
rows, scores = self._index.search(q, min(self._n, max(k * 3, 24)),
beam=self.cfg.beam_width)
out = [(self._ids[int(r)], float(s)) for r, s in zip(rows, scores)]
if candidate_ids is not None:
out = [(fid, s) for fid, s in out if fid in candidate_ids]
out.sort(key=lambda t: -t[1])
if len(out) < k and candidate_ids is not None and len(candidate_ids) > 0:
extra = self._flat_search(q, k, candidate_ids, exclude={f for f, _ in out})
out.extend(extra)
out.sort(key=lambda t: -t[1])
return out[:k]
self.searches_flat += 1
if candidate_ids is not None:
return self._flat_search(q, k, candidate_ids)
sc = self._score_all(q)
k2 = min(k, self._n)
idx = np.argpartition(-sc, k2 - 1)[:k2]
idx = idx[np.argsort(-sc[idx])]
return [(self._ids[int(i)], float(sc[i])) for i in idx]
def _score_all(self, qv: np.ndarray) -> np.ndarray:
# Rust fast path (Task 6-simd): when cortexm_core is built, route
# the per-row scoring through `batch_dot_i8` (int8 codec) or
# `batch_dot` (fp32 path) instead of the codec's numpy `scores`.
# The Rust kernels are AVX-512 → AVX2+FMA → NEON → scalar and
# process the whole batch in one Python→Rust boundary crossing.
# NumPy remains the exact reference; bit parity ≤1e-5 asserted
# by `tests/test_rust_accel.py::TestSimdKernels::test_batch_dot_*`.
try:
from cortexm import accel
except Exception: # pragma: no cover
accel = None
rust_ok = (accel is not None and accel.RUST_AVAILABLE
and accel.RUST_ENABLED)
n = self._n
d = self.dims
if rust_ok and n > 0 and self._packed is not None:
qf32 = np.ascontiguousarray(qv, dtype=np.float32)
if self.codec.uses_aux and self._aux is not None:
# Int8Codec: Rust returns raw int8·f32 dot products;
# multiply by per-row aux scales to match the codec's
# `scores` semantics (`packed.astype(f32) * aux @ q`).
# Pass the int8 buffer as a flat 1-D view — pyo3 binds
# to `PyReadonlyArray1<i8>` zero-copy.
packed_flat = np.ascontiguousarray(
self._packed[:n]).reshape(-1)
raw = np.asarray(
accel._core.batch_dot_i8(packed_flat, qf32, n, d),
dtype=np.float32)
aux = np.asarray(self._aux[:n], dtype=np.float32)
return raw * aux
if self._packed.dtype == np.float32:
# FP32-packed codec path. (BinaryCodec stores bipolar
# ±1 in uint8 — skip; codec.scores handles its own XOR.)
rows_flat = np.ascontiguousarray(
self._packed[:n]).reshape(-1)
raw = np.asarray(
accel._core.batch_dot(rows_flat, qf32, n, d),
dtype=np.float32)
return raw
if self.codec.uses_aux:
return self.codec.scores(self._packed[: self._n], qv, self._aux[: self._n])
return self.codec.scores(self._packed[: self._n], qv)
def _flat_search(self, q: np.ndarray, k: int, candidate_ids: set[str],
exclude: set[str] | None = None) -> list[tuple[str, float]]:
# iterate candidates in PALACE ROW ORDER (insertion order), never in
# set order: set iteration is hash-randomized per process, and the
# argsort below breaks score ties by position — random positions
# would make identical runs return different facts.
sel = [self._id2row[f] for f in candidate_ids
if f in self._id2row
and (exclude is None or f not in exclude)]
sel.sort()
if not sel:
return []
arr = np.array(sel, dtype=np.int64)
packed = self._packed[arr]
aux = self._aux[arr] if self._aux is not None else None
sc = (self.codec.scores(packed, q, aux) if self.codec.uses_aux
else self.codec.scores(packed, q))
order = np.argsort(-sc, kind="stable")[:k]
return [(self._ids[int(arr[i])], float(sc[i])) for i in order]
# ------------------------------------------------------ self-healing
def record_hash(self, row: int) -> str:
blob = self._record_bytes(row)
return self.hasher.hash_bytes(blob)
def _record_bytes(self, row: int) -> bytes:
if isinstance(self.codec, Int8Codec):
return self.codec.to_bytes(self._packed[row], self._aux[row])
return self.codec.to_bytes(self._packed[row])
def stored_hash(self, fact_id: str) -> str | None:
r = self.store.conn.execute(
"SELECT vec_hash FROM vectors WHERE fact_id=?", (fact_id,)).fetchone()
return r["vec_hash"] if r else None
def corrupt(self, rate: float, seed: int = 0,
persist: bool = False) -> int:
"""Inject bit flips (cosmic rays / flash degradation). Returns count."""
if self._n == 0:
return 0
rng = np.random.default_rng(seed)
count = 0
for row in range(self._n):
if rng.random() < rate:
packed_row = self._packed[row]
self._packed[row] = self.codec.corrupt(packed_row, rate, rng)
count += 1
if persist:
blob = self._record_bytes(row)
self.store.conn.execute(
"UPDATE vectors SET record=? WHERE fact_id=?",
(blob, self._ids[row]))
if persist:
self.store.conn.commit()
self._index_valid = False
return count
def health_check(self, sample: int | None = None) -> dict:
"""Detect corrupt records via stored vec_hash; TMR adds bit-level votes."""
import random as _random
rng = _random.Random(7)
rows = list(range(self._n))
if sample and sample < self._n:
rows = rng.sample(rows, sample)
corrupt_rows, checked = [], 0
tmr_flips = 0
for row in rows:
checked += 1
fid = self._ids[row]
want = self.stored_hash(fid)
got = self.record_hash(row)
if want and want != got:
corrupt_rows.append(fid)
if self.cfg.tmr and self.codec.name == "binary":
tmr_flips += self.codec.tmr_health(self._packed[row])
return {"checked": checked, "corrupt": len(corrupt_rows),
"corrupt_ids": corrupt_rows, "tmr_disagree_bits": int(tmr_flips)}
def heal(self, facts_by_id: dict[str, Fact]) -> dict:
"""Re-encode corrupt records from the symbolic Trace (source of truth)."""
report = self.health_check()
healed = 0
for fid in report["corrupt_ids"]:
fact = facts_by_id.get(fid)
if fact is None:
continue
row = self._id2row.get(fid)
if row is None:
continue
vec = self.encode_fact(fact)
packed_row, scale = self._encode_row(vec)
if isinstance(self.codec, Int8Codec):
self._packed[row] = packed_row
self._aux[row] = scale
else:
self._packed[row] = packed_row
blob = (self.codec.to_bytes(packed_row, scale)
if isinstance(self.codec, Int8Codec)
else self.codec.to_bytes(packed_row))
self.store.conn.execute(
"UPDATE vectors SET record=?, vec_hash=? WHERE fact_id=?",
(blob, self.hasher.hash_bytes(blob), fid))
healed += 1
self.store.conn.commit()
self._index_valid = False
return {"corrupt": report["corrupt"], "healed": healed,
"note": "re-encoded from Trace; source hash verified"}
# ------------------------------------------------------------- stats
def storage_stats(self) -> dict:
vec_bytes = self._n * self.codec.bytes_per_vector if self._n else 0
return {
"vectors": self._n,
"codec": self.codec.name,
"bytes_per_vector": self.codec.bytes_per_vector,
"vector_bytes": int(vec_bytes),
"per_million_memories_mb": round(
self.codec.bytes_per_vector * 1e6 / 1e6, 1),
"index": ("tree" if self._index_valid else
("flat" if self._n else "empty")),
"searches_flat": self.searches_flat,
"searches_indexed": self.searches_indexed,
}