SPIRA / python /spira /ept.py
thefinalboss's picture
SPIRA v2: entropy-patch tokenization + hierarchical dopamine-gated spiking cortex + hippocampal one-shot memory with replay consolidation
032a961 verified
Raw History Blame Contribute Delete
5.87 kB
"""EPT — Entropy-Patch Tokenization.
Segments a byte stream into variable-length patches such that each patch
carries approximately ``budget_bits`` of accumulated surprisal under an
order-1 byte-level Markov model ("constant-information quanta").
Consequences:
* low-entropy regions (repeated keywords, indentation, boilerplate) produce
LONG patches, high-entropy regions produce SHORT ones — no merge table,
no vocabulary to train;
* patches are content-addressed with FNV-1a-64: the same byte substring is
always the same token, wherever it occurs, with zero learning;
* strictly streaming: one pass, O(n) with tiny state.
"""
from __future__ import annotations
import numpy as np
from .hashing import fnv1a64
class EPT:
def __init__(
self,
budget_bits: float = 16.0,
min_len: int = 2,
max_len: int = 32,
alpha: float = 0.5,
) -> None:
self.budget_bits = float(budget_bits)
self.min_len = int(min_len)
self.max_len = int(max_len)
self.alpha = float(alpha)
self.log2p: np.ndarray | None = None # (256, 256) float32
# ------------------------------------------------------------------ fit
def fit(self, data: bytes) -> "EPT":
"""Estimate the order-1 byte Markov model on training bytes."""
arr = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
C = np.zeros((256, 256), dtype=np.int64)
np.add.at(C, (arr[:-1], arr[1:]), 1)
P = (C + self.alpha).astype(np.float64)
P /= P.sum(axis=1, keepdims=True)
self.log2p = np.log2(P).astype(np.float32)
return self
# ------------------------------------------------------------ surprise
def surprise(self, data: bytes) -> np.ndarray:
"""Per-byte surprisal -log2 P(byte | previous byte), in bits."""
if self.log2p is None:
raise RuntimeError("EPT must be fit() before use")
arr = np.frombuffer(data, dtype=np.uint8).astype(np.int64)
s = np.empty(len(arr), dtype=np.float32)
if len(arr) == 0:
return s
s[0] = 8.0 # uniform prior for the first byte
if len(arr) > 1:
s[1:] = -self.log2p[arr[:-1], arr[1:]]
return s
# ----------------------------------------------------------- segmenting
def boundaries(self, data: bytes) -> np.ndarray:
"""Greedy constant-surprise segmentation.
Returns inclusive boundary positions ``b_0=0 < b_1 < ... < b_k = n``;
patch i is ``data[b_i : b_{i+1}]``. A patch closes when its
accumulated surprisal reaches ``budget_bits`` (or ``max_len`` /
end-of-stream forces it).
"""
if len(data) == 0:
return np.zeros(0, dtype=np.int64)
s = self.surprise(data)
n = len(s)
cum = np.concatenate([[0.0], np.cumsum(s, dtype=np.float64)])
bounds = [0]
start = 0
searchsorted = np.searchsorted
min_len, max_len, budget = self.min_len, self.max_len, self.budget_bits
while start < n:
target = cum[start] + budget
b = int(searchsorted(cum, target, side="left"))
if b < start + min_len:
b = start + min_len
elif b > start + max_len:
b = start + max_len
if b > n:
b = n
bounds.append(b)
start = b
# drop a degenerate trailing boundary at exactly n produced twice
if len(bounds) > 1 and bounds[-1] == bounds[-2]:
bounds.pop()
return np.asarray(bounds, dtype=np.int64)
def segment(self, data: bytes) -> list[bytes]:
b = self.boundaries(data)
return [data[b[i] : b[i + 1]] for i in range(len(b) - 1)]
def segment_ids(self, data: bytes) -> list[int]:
return [fnv1a64(p) for p in self.segment(data)]
# ---------------------------------------------------------- calibration
def calibrate(self, sample: bytes, target_len: float = 4.5, iters: int = 10) -> float:
"""Binary-search the surprisal budget to hit a target average patch
length on a calibration sample. Returns the chosen budget."""
lo, hi = 1.0, 64.0
for _ in range(iters):
mid = 0.5 * (lo + hi)
self.budget_bits = mid
b = self.boundaries(sample)
avg = len(sample) / max(len(b) - 1, 1)
if avg < target_len:
lo = mid
else:
hi = mid
self.budget_bits = 0.5 * (lo + hi)
return self.budget_bits
# ------------------------------------------------------------- metrics
def stats(self, patches: list[bytes]) -> dict:
lens = np.array([len(p) for p in patches], dtype=np.float64)
return {
"n_patches": int(len(patches)),
"avg_len": float(lens.mean()) if len(lens) else 0.0,
"min_len": int(lens.min()) if len(lens) else 0,
"max_len": int(lens.max()) if len(lens) else 0,
"p50_len": float(np.percentile(lens, 50)) if len(lens) else 0.0,
"p90_len": float(np.percentile(lens, 90)) if len(lens) else 0.0,
"bytes": int(lens.sum()),
}
# ---------------------------------------------------------- (de)hydrate
def state_dict(self) -> dict:
return {
"budget_bits": self.budget_bits,
"min_len": self.min_len,
"max_len": self.max_len,
"alpha": self.alpha,
"log2p": self.log2p.astype(np.float32),
}
def load_state_dict(self, sd: dict) -> None:
self.budget_bits = float(sd["budget_bits"])
self.min_len = int(sd["min_len"])
self.max_len = int(sd["max_len"])
self.alpha = float(sd["alpha"])
self.log2p = np.asarray(sd["log2p"], dtype=np.float32)