File size: 5,872 Bytes
032a961 | 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """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)
|