File size: 4,394 Bytes
5d97275 | 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 | """Batched, vectorised byte n-gram hashing (numpy) - the shared input of the tier-1 HTTP model at training time
(torch EmbeddingBag) and inference time (numpy). Same function both sides, so the buckets always agree.
A batch of texts becomes one uint8 array with a 0 byte between texts; for each n-gram order the rolling polynomial
hash of every window is computed with a few whole-array operations, windows that cross a text boundary are dropped,
and a 64-bit mix maps the hash to a bucket. No per-character python, so cost is dominated by memory bandwidth.
"""
from __future__ import annotations
import numpy as np
P = np.uint64(0x100000001B3) # FNV-64 prime
SEED = np.uint64(0x9E3779B97F4A7C15)
M1 = np.uint64(0xBF58476D1CE4E5B9) # splitmix64 finaliser
M2 = np.uint64(0x94D049BB133111EB)
ORDERS = (2, 3, 4, 5, 6)
with np.errstate(over="ignore"):
ORDER_SEEDS = {n: np.uint64(SEED * np.uint64(n)) for n in range(1, 17)} # wraps mod 2^64 by design
def _mix(h: np.ndarray) -> np.ndarray:
h = h ^ (h >> np.uint64(30))
h = h * M1
h = h ^ (h >> np.uint64(27))
h = h * M2
return h ^ (h >> np.uint64(31))
def pack(texts: list, max_bytes: int = 1024) -> tuple[np.ndarray, np.ndarray]:
"""texts -> (bytes with a 0 separator after each text, start offset of each text). Texts are cut at max_bytes."""
enc = [(t if isinstance(t, (bytes, bytearray)) else str(t).encode("utf-8", "replace"))[:max_bytes].replace(b"\0", b" ")
for t in texts]
buf = np.frombuffer(b"\0".join(enc) + b"\0", dtype=np.uint8)
lens = np.fromiter((len(e) for e in enc), dtype=np.int64, count=len(enc))
starts = np.concatenate(([0], np.cumsum(lens + 1)[:-1])) if len(enc) else np.zeros(0, np.int64)
return buf, starts
@np.errstate(over="ignore")
def ngram_buckets(texts: list, buckets: int, orders=ORDERS, max_bytes: int = 1024) -> tuple[np.ndarray, np.ndarray]:
"""-> (bucket index per n-gram, text index per n-gram), both int64, grouped by order then position.
Texts shorter than the smallest order contribute nothing (callers treat an empty bag as all-zero)."""
buf, starts = pack(texts, max_bytes)
n_all = buf.size
b64 = buf.astype(np.uint64)
sep = (buf == 0).astype(np.int32)
sep_cum = np.concatenate(([0], np.cumsum(sep))) # separators before position i
doc_of = np.searchsorted(starts, np.arange(n_all), side="right") - 1
h = b64 * P
out_idx, out_doc = [], []
for n in range(1, max(orders) + 1):
if n > 1:
h = h[:-1] * P + b64[n - 1:] # hash of buf[i:i+n]
if n not in orders:
continue
m = n_all - n + 1
ok = (sep_cum[n:n + m] - sep_cum[:m]) == 0 # no separator inside the window
hv = _mix(h[:m][ok] ^ ORDER_SEEDS[n])
out_idx.append((hv % np.uint64(buckets)).astype(np.int64))
out_doc.append(doc_of[:m][ok])
if not out_idx:
return np.zeros(0, np.int64), np.zeros(0, np.int64)
return np.concatenate(out_idx), np.concatenate(out_doc)
def bag_mean(emb: np.ndarray, idx: np.ndarray, doc: np.ndarray, n_docs: int, scale: float | np.ndarray = 1.0) -> np.ndarray:
"""Mean of embedding rows per text (EmbeddingBag mode='mean'); emb may be int8 with a per-column `scale`."""
d = emb.shape[1]
out = np.zeros((n_docs, d), np.float32)
if idx.size:
# ngram_buckets output is sorted by text inside each order block: segment sums with reduceat (C speed),
# instead of np.add.at (~5x slower)
cuts = np.concatenate(([0], np.flatnonzero(np.diff(doc) < 0) + 1, [doc.size]))
for s, e in zip(cuts[:-1], cuts[1:]):
db = doc[s:e]
first = np.flatnonzero(np.concatenate(([True], db[1:] != db[:-1])))
out[db[first]] += np.add.reduceat(emb[idx[s:e]], first, axis=0, dtype=np.float32)
cnt = np.bincount(doc, minlength=n_docs).astype(np.float32)
out /= np.maximum(cnt, 1)[:, None]
return out * np.asarray(scale, np.float32)
def offsets_for_torch(doc: np.ndarray, n_docs: int) -> tuple[np.ndarray, np.ndarray]:
"""(order of n-grams sorted by text, offsets) for torch.nn.EmbeddingBag(input, offsets)."""
order = np.argsort(doc, kind="stable")
counts = np.bincount(doc, minlength=n_docs)
return order, np.concatenate(([0], np.cumsum(counts)[:-1]))
|