| """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) |
| SEED = np.uint64(0x9E3779B97F4A7C15) |
| M1 = np.uint64(0xBF58476D1CE4E5B9) |
| 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)} |
|
|
|
|
| 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))) |
| 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:] |
| if n not in orders: |
| continue |
| m = n_all - n + 1 |
| ok = (sep_cum[n:n + m] - sep_cum[:m]) == 0 |
| 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: |
| |
| |
| 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])) |
|
|
