cortexm 0.3.0__py3-none-any.whl

cortexm/vsa/role_vectors.py

@@ -0,0 +1,236 @@
+"""Engineered role vectors — NSR-inspired (ESWEEK24).
+
+arXiv insight: NSR trains an autoencoder to learn compact embeddings
+for the fact vocabulary, then uses those embeddings to construct role
+vectors that are *maximally mutually orthogonal* in the data manifold.
+Random role vectors work (HDC theory guarantees capacity grows with
+sqrt(dims)), but engineered ones have higher effective capacity and
+lower cross-talk because they sit on the directions of greatest
+variance in the actual data.
+
+This module provides:
+
+    EngineeredRoleVectors(dims, n_roles=3)
+        .fit(fact_matrix)        -> trains a tiny 1-layer autoencoder
+                                     on the fact matrix and stores the
+                                     top-k principal directions as role
+                                     vectors
+        .role_vec(role)          -> returns the engineered role vector
+                                     (falls back to VSA's random init
+                                     if not yet fit)
+        .save(path) / .load(path)
+
+WHY THIS MATTERS:
+  - role_vec("S") currently = rng.standard_normal(dims). For dims=768
+    that's a random point on the unit sphere — fine in theory but it
+    wastes capacity on directions orthogonal to the actual fact vocab.
+  - After fit(), role_vec("S") = the top-1 principal direction of the
+    S-vector manifold. role_vec("R") = top-2, role_vec("V") = top-3.
+    These directions are where the data actually lives, so:
+      (a) bound holograms exploit the full effective capacity
+      (b) cross-talk between S/R/V role bindings drops because the
+          top-k principal directions are approximately orthogonal by
+          construction
+      (c) retrieval probes have higher SNR because they're aligned
+          with the data axes
+
+IMPLEMENTATION:
+  Tiny 1-layer linear autoencoder (no nonlinearity) trained with
+  plain SGD on the fact matrix. The encoder weights converge to the
+  top-k principal components (proven by Baldi & Horn, 1989 — linear
+  AE on centered data recovers PCA). We center the data and train the
+  encoder to have orthogonal rows via a soft orthogonality penalty.
+
+  Cost: ~1 day of human work, ~30 seconds of CPU time on 1000 facts.
+  Deterministic given the seed.
+"""
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import numpy as np
+
+
+class EngineeredRoleVectors:
+    """Train and serve engineered role vectors.
+
+    Construction:
+        erv = EngineeredRoleVectors(dims=768, n_roles=3, seed=42)
+
+    Fit:
+        erv.fit(fact_matrix)  # fact_matrix: (n_facts, dims)
+        # trains a tiny AE, extracts top-k principal directions,
+        # stores them as the role vectors
+
+    Use:
+        erv.role_vec("S")  # top-1 principal direction
+        erv.role_vec("R")  # top-2
+        erv.role_vec("V")  # top-3
+    """
+
+    ROLE_ORDER = ["S", "R", "V"]  # subject / relation / value
+
+    def __init__(self, dims: int = 768, n_roles: int = 3,
+                 seed: int = 42, n_epochs: int = 200,
+                 lr: float = 0.01, orth_penalty: float = 0.1) -> None:
+        self.dims = dims
+        self.n_roles = n_roles
+        self.seed = seed
+        self.n_epochs = n_epochs
+        self.lr = lr
+        self.orth_penalty = orth_penalty
+        self._role_vecs: dict[str, np.ndarray] = {}
+        self._mean: np.ndarray | None = None
+        self._fit_loss: list[float] = []
+
+    # ------------------------------------------------------------- fit
+    def fit(self, fact_matrix: np.ndarray) -> dict:
+        """Train a tiny linear autoencoder on `fact_matrix`.
+
+        fact_matrix: (n_samples, dims) float32. Should contain the
+            subject / relation / value vectors for all facts in the
+            corpus, stacked. The AE learns to reconstruct them through
+            a k-dim bottleneck where k = n_roles.
+
+        After training, the encoder's rows are the top-k principal
+        directions. We store them as role_vec("S"/"R"/"V").
+
+        Returns a training report dict.
+        """
+        X = np.asarray(fact_matrix, dtype=np.float32)
+        if X.ndim != 2:
+            raise ValueError(f"expected 2D matrix, got shape {X.shape}")
+        n, d = X.shape
+        if d != self.dims:
+            raise ValueError(
+                f"matrix dim {d} != configured dims {self.dims}")
+        if n < self.n_roles:
+            # not enough samples to learn k directions — fall back
+            # to top-k random vectors orthogonalized via Gram-Schmidt
+            rng = np.random.default_rng(self.seed)
+            vecs = rng.standard_normal((self.n_roles, d)).astype(np.float32)
+            for i in range(self.n_roles):
+                for j in range(i):
+                    vecs[i] -= np.dot(vecs[i], vecs[j]) * vecs[j]
+                n_ = float(np.linalg.norm(vecs[i]))
+                vecs[i] /= max(n_, 1e-9)
+            for i, r in enumerate(self.ROLE_ORDER[:self.n_roles]):
+                self._role_vecs[r] = vecs[i]
+            return {"trained": False, "reason": "insufficient_samples",
+                    "n_samples": n, "n_roles": self.n_roles,
+                    "fallback": "gram_schmidt_random"}
+
+        # center the data (PCA assumes centered data)
+        self._mean = X.mean(axis=0)
+        Xc = X - self._mean
+
+        # tiny linear AE: encoder W: (k, d), decoder W.T: (d, k)
+        # init with small random values
+        rng = np.random.default_rng(self.seed)
+        k = self.n_roles
+        scale = float(1.0 / np.sqrt(d))
+        W = rng.standard_normal((k, d)).astype(np.float32) * scale
+
+        # SGD on reconstruction loss + orthogonality penalty
+        # (orth penalty: W @ W.T should be ≈ I)
+        batch_size = min(64, n)
+        for epoch in range(self.n_epochs):
+            perm = rng.permutation(n)
+            epoch_loss = 0.0
+            for i in range(0, n, batch_size):
+                idx = perm[i:i + batch_size]
+                xb = Xc[idx]                      # (b, d)
+                # forward: z = xb @ W.T   (b, k)
+                # recon: xhat = z @ W     (b, d)
+                z = xb @ W.T
+                xhat = z @ W
+                # L2 reconstruction loss per sample — clip to avoid
+                # overflow on numerically large data (the orth penalty
+                # can push weights to blow up if lr is too high)
+                diff = xhat - xb
+                # clip per-element to a safe range
+                diff = np.clip(diff, -1e4, 1e4)
+                loss = float(np.mean(np.sum(diff * diff, axis=1)))
+                # grad on W: dL/dW = 2 * (xhat - x).T @ z / b
+                # shape (d, k) -> transpose for our layout
+                g = 2.0 * diff.T @ z / len(idx)   # (d, k)
+                # clip gradient to prevent overflow → NaN
+                g = np.clip(g, -1.0, 1.0)
+                # orthogonality penalty: ||W @ W.T - I||_F^2 / k
+                # grad: 2/k * (W @ W.T - I) @ W
+                WtW = W @ W.T                      # (k, k)
+                I_k = np.eye(k, dtype=np.float32)
+                orth_grad = (2.0 / k) * (WtW - I_k) @ W
+                orth_grad = np.clip(orth_grad, -1.0, 1.0)
+                # combined grad on W (note: g is (d,k), so transpose)
+                W -= self.lr * (g.T + self.orth_penalty * orth_grad)
+                # also clip weights themselves to a safe range
+                W = np.clip(W, -10.0, 10.0)
+                epoch_loss += loss * len(idx)
+            epoch_loss /= n
+            self._fit_loss.append(epoch_loss)
+            if epoch % 20 == 0 or epoch == self.n_epochs - 1:
+                # report conditioning of W @ W.T (lower = more orthogonal)
+                cond = float(np.linalg.cond(WtW)) if k > 1 else 1.0
+                # suppress per-epoch logging during normal use
+                pass
+
+        # extract role vectors: rows of W are the top-k principal dirs
+        # normalize to unit length
+        for i, r in enumerate(self.ROLE_ORDER[:k]):
+            v = W[i].copy()
+            nrm = float(np.linalg.norm(v))
+            self._role_vecs[r] = v / max(nrm, 1e-9)
+
+        # report final reconstruction loss + orthogonality
+        WtW_final = W @ W.T
+        off_diag = float(np.sum(np.abs(WtW_final - np.eye(k, dtype=np.float32))))
+        return {
+            "trained": True,
+            "n_samples": n,
+            "dims": d,
+            "n_roles": k,
+            "epochs": self.n_epochs,
+            "final_loss": self._fit_loss[-1] if self._fit_loss else 0.0,
+            "initial_loss": self._fit_loss[0] if self._fit_loss else 0.0,
+            "loss_reduction_pct": (
+                (1.0 - self._fit_loss[-1] / max(self._fit_loss[0], 1e-9)) * 100
+                if self._fit_loss else 0.0),
+            "off_diag_sum": off_diag,
+            "condition_number": (float(np.linalg.cond(WtW_final))
+                                  if k > 1 else 1.0),
+        }
+
+    # ------------------------------------------------------------- serve
+    def role_vec(self, role: str) -> np.ndarray | None:
+        return self._role_vecs.get(role)
+
+    @property
+    def is_fit(self) -> bool:
+        return bool(self._role_vecs)
+
+    def save(self, path: str | os.PathLike) -> None:
+        """Persist the role vectors to a .npz file."""
+        arrays = {f"role_{r}": v for r, v in self._role_vecs.items()}
+        if self._mean is not None:
+            arrays["mean"] = self._mean
+        arrays["meta"] = np.array(json.dumps({
+            "dims": self.dims, "n_roles": self.n_roles,
+            "seed": self.seed, "final_loss": (
+                self._fit_loss[-1] if self._fit_loss else 0.0),
+        }), dtype=str)
+        np.savez(path, **arrays)
+
+    def load(self, path: str | os.PathLike) -> None:
+        data = np.load(path, allow_pickle=False)
+        for r in self.ROLE_ORDER[:self.n_roles]:
+            key = f"role_{r}"
+            if key in data:
+                self._role_vecs[r] = data[key]
+        if "mean" in data:
+            self._mean = data["mean"]
+
+
+__all__ = ["EngineeredRoleVectors"]

cortexm/vsa/slb.py

@@ -0,0 +1,78 @@
+"""Semantic Lookaside Buffer — conversational locality cache.
+
+64-entry ring buffer of quantized query signatures with cached result
+sets. A new query whose signature is ≥ ``threshold`` cosine-similar to a
+cached signature reuses the cached ranking (conversational locality:
+follow-up questions are near-duplicates of their predecessors). L1-
+resident by design; hit path costs one 64×dims dot product.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+class SemanticLookasideBuffer:
+    def __init__(self, entries: int = 64, threshold: float = 0.97,
+                 dims: int = 768) -> None:
+        self.capacity = entries
+        self.threshold = threshold
+        self.dims = dims
+        self._sigs = np.zeros((entries, dims), dtype=np.float32)
+        self._results: list[list[tuple[str, float]] | None] = [None] * entries
+        self._queries: list[str | None] = [None] * entries
+        self._scopes: list[tuple | None] = [None] * entries
+        self._pos = 0
+        self._filled = 0
+        self.hits = 0
+        self.misses = 0
+        self.total_hit_latency = 0.0
+        self.total_miss_latency = 0.0
+
+    def lookup(self, q: np.ndarray,
+               scope: tuple | None = None) -> list[tuple[str, float]] | None:
+        """Return cached results for a signature **in the same scope**.
+
+        Scope-blind lookup is a correctness bug, not just a privacy one:
+        near-duplicate queries from different users would cross-contaminate
+        (and then die in the caller's scope filter, yielding empty blocks).
+        """
+        if self._filled == 0:
+            return None
+        sims = self._sigs[: self._filled] @ q
+        best = int(np.argmax(sims))
+        if float(sims[best]) >= self.threshold and self._scopes[best] == scope:
+            self.hits += 1
+            return self._results[best]
+        return None
+
+    def store(self, q: np.ndarray, results: list[tuple[str, float]],
+              query: str = "", scope: tuple | None = None) -> None:
+        pos = self._pos
+        self._sigs[pos] = q
+        self._results[pos] = results
+        self._queries[pos] = query
+        self._scopes[pos] = scope
+        self._pos = (self._pos + 1) % self.capacity
+        self._filled = min(self._filled + 1, self.capacity)
+
+    def record_latency(self, hit: bool, seconds: float) -> None:
+        if hit:
+            self.total_hit_latency += seconds
+        else:
+            self.total_miss_latency += seconds
+
+    @property
+    def miss_latency_avg(self) -> float:
+        return (self.total_miss_latency / self.misses) if self.misses else 0.0
+
+    def stats(self) -> dict:
+        total = self.hits + self.misses
+        return {
+            "hits": self.hits, "misses": self.misses,
+            "hit_rate": round(self.hits / total, 4) if total else 0.0,
+            "avg_hit_latency_us": round(
+                self.total_hit_latency / self.hits * 1e6, 1) if self.hits else 0.0,
+            "avg_miss_latency_us": round(self.miss_latency_avg * 1e6, 1),
+            "entries_used": self._filled,
+        }

cortexm/vsa/tlsh_trie.py

@@ -0,0 +1,137 @@
+"""TLSH ternary trie — software TCAM for O(log N + w) hologram lookup.
+
+The Stanford TLSH paper (arXiv:1006.3514) uses ternary content-
+addressable memory for O(1) parallel lookup. Without TCAM hardware,
+we emulate it as a ternary Patricia trie over packed binary holograms:
+each path is the bits of a packed binary vector, with wildcard edges
+that match either bit (the ternary '*' bit).
+
+Use as a *pre-filter* in MemoryPalace.search when codec is binary/rabitq:
+the trie returns O(k) candidate fact_ids within max_wildcards bit-flips
+of the query, then codec.scores ranks them. Trades one big argsort for
+a sub-linear trie walk — wins when N grows large and dims is high (16k+).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+class _TrieNode:
+    __slots__ = ("children", "fact_ids", "depth")
+
+    def __init__(self, depth: int) -> None:
+        self.children: dict[int, _TrieNode] = {}  # bit 0/1 → node
+        self.fact_ids: list[str] = []
+        self.depth = depth
+
+
+class TernaryTrie:
+    """Software TLSH over packed binary holograms."""
+
+    def __init__(self, dims: int, max_wildcards: int = 8,
+                 max_candidates: int = 256) -> None:
+        self.dims = dims
+        self.max_wildcards = max_wildcards
+        self.max_candidates = max_candidates
+        self.root = _TrieNode(0)
+        self._size = 0
+
+    def insert(self, fact_id: str, packed: np.ndarray) -> None:
+        """Insert a packed binary vector (uint8 array of packed bits)."""
+        bits = _unpack_bits(packed, self.dims)
+        node = self.root
+        for bit in bits:
+            b = int(bit)
+            child = node.children.get(b)
+            if child is None:
+                child = _TrieNode(node.depth + 1)
+                node.children[b] = child
+            node = child
+        node.fact_ids.append(fact_id)
+        self._size += 1
+
+    def lookup(self, packed_q: np.ndarray, k: int = 10,
+               max_wildcards: int | None = None) -> list[tuple[str, int]]:
+        """Return up to k (fact_id, hamming_distance) pairs within
+        max_wildcards bit-flips of packed_q. O(log N + max_wildcards·branch).
+        """
+        mw = max_wildcards if max_wildcards is not None else self.max_wildcards
+        bits = _unpack_bits(packed_q, self.dims)
+        # candidate (node, position, wildcards_used, distance_so_far)
+        # use a stack with priority by wildcards_used
+        results: list[tuple[str, int]] = []
+        seen: set[str] = set()
+        # DFS with wildcard budget
+        # stack entries: (node, idx, wildcards_remaining)
+        # we don't strictly bound exploration — for small max_wildcards
+        # and a sparse trie this is fine
+        stack = [(self.root, 0, mw)]
+        # use a heap for best-first with priority on wildcards remaining
+        import heapq
+        # priority: (-wildcards_remaining, idx) so most wildcards remaining
+        # (i.e. least used) comes first
+        heap: list[tuple[int, int, _TrieNode]] = [(-mw, 0, self.root)]
+        while heap and len(results) < self.max_candidates:
+            _, idx, node = heapq.heappop(heap)
+            if node.fact_ids and idx == self.dims:
+                for fid in node.fact_ids:
+                    if fid not in seen:
+                        seen.add(fid)
+                        results.append((fid, mw + _heap_key(0, mw)))
+                        if len(results) >= self.max_candidates:
+                            break
+                continue
+            if idx >= self.dims:
+                # at a leaf but bits remaining — these are stored facts
+                for fid in node.fact_ids:
+                    if fid not in seen:
+                        seen.add(fid)
+                        results.append((fid, mw))
+                continue
+            want_bit = int(bits[idx])
+            # exact match: free (no wildcard used)
+            exact = node.children.get(want_bit)
+            if exact is not None:
+                heapq.heappush(heap, (-(mw), idx + 1, exact))
+            # wildcard match: try the other bit (uses 1 wildcard)
+            other = 1 - want_bit
+            wild = node.children.get(other)
+            if wild is not None and mw > 0:
+                heapq.heappush(heap, (-(mw - 1), idx + 1, wild))
+        # dedupe and sort by best distance estimate
+        # NB: distance estimate is approximate (lower bound by wildcards used)
+        dedup: dict[str, int] = {}
+        for fid, dist in results:
+            if fid not in dedup or dist < dedup[fid]:
+                dedup[fid] = dist
+        out = sorted(dedup.items(), key=lambda x: x[1])[:k]
+        return out
+
+    def __len__(self) -> int:
+        return self._size
+
+    def stats(self) -> dict:
+        return {
+            "dims": self.dims,
+            "size": self._size,
+            "max_wildcards": self.max_wildcards,
+            "root_children": len(self.root.children),
+        }
+
+
+def _heap_key(used: int, budget: int) -> int:
+    """Convert 'wildcards used' into a priority value (less used = better)."""
+    return budget - used  # remaining
+
+
+def _unpack_bits(packed: np.ndarray, dims: int) -> np.ndarray:
+    """Unpack packed uint8 bits into a 1D array of 0/1 of length dims."""
+    arr = np.atleast_1d(packed)
+    if arr.dtype == np.uint8 and len(arr) * 8 >= dims:
+        return np.unpackbits(arr, count=dims).astype(np.uint8)
+    # already a bit array
+    return arr.astype(np.uint8)
+
+
+__all__ = ["TernaryTrie"]