File size: 5,088 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 | """Baselines for comparison against SPIRA.
Compression: gzip -9, HuggingFace byte-level BPE (matched vocabulary sizes).
Prediction: unigram, patch-bigram with unigram backoff.
"""
from __future__ import annotations
import gzip
import time
import numpy as np
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
# ------------------------------------------------------------- compression
def gzip_bits_per_byte(data: bytes) -> tuple[float, float]:
"""(bits/byte, seconds) for gzip level 9."""
t0 = time.perf_counter()
comp = gzip.compress(data, 9)
dt = time.perf_counter() - t0
return len(comp) * 8 / len(data), dt
def bpe_encode(train: bytes, test: bytes, vocab_size: int) -> dict:
"""Train a byte-level BPE on `train` and encode `test` with it.
Returns dict(ids, vocab_size, vocab_bytes, train_s, encode_s,
MB_per_s_train).
"""
t0 = time.perf_counter()
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
special_tokens=[],
)
tok.train_from_iterator([train.decode("utf-8", "replace")], trainer=trainer)
train_s = time.perf_counter() - t0
t1 = time.perf_counter()
enc = tok.encode(test.decode("utf-8", "replace"))
encode_s = time.perf_counter() - t1
vocab = tok.get_vocab()
vocab_bytes = sum(len(s.encode("utf-8")) for s in vocab)
return {
"ids": np.asarray(enc.ids, dtype=np.int64),
"vocab_size": len(vocab),
"vocab_bytes": vocab_bytes,
"train_s": train_s,
"encode_s": encode_s,
}
def stream_entropy_bits(ids: np.ndarray, order: int = 0) -> float:
"""Empirical entropy of a token stream in bits/token (order 0 or 1)."""
# remap arbitrary ids (e.g. uint64 patch ids) to dense indices
uniq, inv = np.unique(ids, return_inverse=True)
ids = inv.astype(np.int64)
if order == 0:
_, cnt = np.unique(ids, return_counts=True)
p = cnt / cnt.sum()
return float(-(p * np.log2(p)).sum())
# order 1: H(next | prev) estimated from the same stream
prev = ids[:-1]
cur = ids[1:]
span = len(uniq)
pair = prev * span + cur
uniq_pair, cnt_pair = np.unique(pair, return_counts=True)
pair_prev = uniq_pair // span
prev_ids, cnt_prev = np.unique(prev, return_counts=True)
lookup = np.zeros(span, dtype=np.int64)
lookup[prev_ids] = cnt_prev
denom = lookup[pair_prev].astype(np.float64)
p_pair = cnt_pair / denom
bits = -(cnt_pair * np.log2(p_pair)).sum()
return float(bits / len(cur))
def dict_bits(vocab_size: int, vocab_bytes: int, overhead_bits: float = 32.0) -> float:
"""Decode-dictionary storage: raw token bytes + id/length overhead."""
return vocab_bytes * 8 + vocab_size * overhead_bits
# ------------------------------------------------------------- prediction
class UnigramBaseline:
def __init__(self, train_cols: np.ndarray, V: int | None = None) -> None:
if V is not None:
train_cols = train_cols[train_cols != V] # exclude the UNK column
self.best = int(np.bincount(train_cols).argmax())
def accuracy(self, cur: np.ndarray) -> float:
return float((cur == self.best).mean())
class BigramBaseline:
"""Patch-bigram predictor with unigram backoff (sparse pair counts).
Dense (V, V) storage is infeasible at V=16k (1+ GB); pairs are counted
via packed int64 keys instead, and the argmax continuation per preceding
column is precomputed by a (prev, -count) lexsort.
"""
def __init__(self, train_cols: np.ndarray, V: int) -> None:
prev = train_cols[:-1]
cur = train_cols[1:]
ok = (prev != V) & (cur != V) # in-vocab pairs only
p = prev[ok].astype(np.int64)
c = cur[ok].astype(np.int64)
if len(p) == 0:
self.map_prev = np.array([-1], dtype=np.int64)
self.map_pred = np.array([0], dtype=np.int64)
else:
key = p * np.int64(V + 1) + c
uk, cnt = np.unique(key, return_counts=True) # sorted by key
prev_u = (uk // np.int64(V + 1)).astype(np.int64)
pred_c = (uk % np.int64(V + 1)).astype(np.int64)
order = np.lexsort((-cnt, prev_u)) # prev asc, count desc
pu, pc = prev_u[order], pred_c[order]
first = np.ones(len(pu), dtype=bool)
first[1:] = pu[1:] != pu[:-1]
self.map_prev = pu[first]
self.map_pred = pc[first]
self.uni_best = int(np.bincount(train_cols[train_cols != V]).argmax())
def accuracy(self, prev: np.ndarray, cur: np.ndarray) -> float:
pos = np.searchsorted(self.map_prev, prev)
pos_c = np.minimum(pos, len(self.map_prev) - 1)
hit = (pos < len(self.map_prev)) & (self.map_prev[pos_c] == prev)
pred = np.where(hit, self.map_pred[pos_c], self.uni_best)
return float((pred == cur).mean())
|