HAKO upload: hako/tokenizer/byte_bpe.py

hako/tokenizer/byte_bpe.py

@@ -0,0 +1,234 @@
+"""Parallel byte-level BPE trainer/encoder with provably exact sharded counts.
+
+Lemma 1 (additivity of pair counts across shards -- exactness).
+The corpus is a MULTISET OF DOCUMENT SEQUENCES {s_d}; tokenization is
+document-independent, so the multiset of adjacent pairs decomposes as the
+DISJOINT union
+    Pair(corpus) = U_d Interior(s_d),   Interior(s_d) = {(s_d_i, s_d_{i+1})}.
+A shard assignment partitions the documents, hence for every symbol pair (a,b)
+    C_global(a, b) = sum_shard p  sum_{d in p} C_interior_d(a, b),
+an identity (counting measure over a disjoint union). No cross-document
+pairs exist in the object language (documents are atomic sequences, exactly
+as in LM training), so -- unlike a naive 'boundary pair ownership' scheme --
+NO counting or ownership rule can depend on the sharding. The merge choice
+is a deterministic function of the global counts (argmax, fixed
+lexicographic tie-break, min-frequency threshold), and every occurrence of
+the winning pair lies ENTIRELY inside one document, so the per-shard apply
+step rewrites exactly the occurrences that were counted. By induction on
+merge iterations, a trainer consuming sharded counts performs EXACTLY the
+merge sequence of a full-corpus trainer, for any sharding. Q.E.D.
+
+(Correction note: an earlier draft counted pairs between chained documents
+inside a shard; that made counts sharding-dependent -- the inconsistency is
+resolved by the document-atomic formulation above, which the test suite
+verifies by merge-history equality across shardings.)
+
+Encoding uses the learned rank table (greedy lowest-rank merge); the vocab
+is byte-aligned so any UTF-8 string round-trips losslessly (all 256 bytes
+are always in the vocab: PT-BR accents and CJK included).
+"""
+from __future__ import annotations
+
+import json
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Dict, Iterable, List, Sequence, Tuple
+
+_NUM_WORKERS_CAP = 8
+
+
+def _byte_encode(text: str) -> List[int]:
+    return list(text.encode("utf-8"))
+
+
+def _byte_decode(ids: Sequence[int]) -> str:
+    return bytes(ids).decode("utf-8", errors="replace")
+
+
+class ByteBPETokenizer:
+    """Trainer + encoder/decoder for byte-level BPE with sharded exact counts."""
+
+    def __init__(self, vocab_size: int = 8192, min_freq: int = 2,
+                 shards: int = 8) -> None:
+        self.vocab_size = max(vocab_size, 300)   # TARGET size (attribute)
+        self.min_freq = min_freq
+        self.shards = max(2, min(shards, _NUM_WORKERS_CAP))
+        # id -> bytes string (bytes 0..255 always present)
+        self.vocab: Dict[int, bytes] = {i: bytes([i]) for i in range(256)}
+        self.ranks: Dict[Tuple[int, int], int] = {}  # pair -> merge priority
+        self.merge_left: Dict[Tuple[int, int], int] = {}
+        self.merge_right: Dict[Tuple[int, int], int] = {}
+        self.merge_history: List[Tuple[int, int, int]] = []
+
+    # ------------------------------------------------------------- training
+    @staticmethod
+    def _count_interior(seq: List[int]) -> Counter:
+        c = Counter()
+        for a, b in zip(seq, seq[1:]):
+            c[(a, b)] += 1
+        return c
+
+    def _boundary_counts(self, shard_seqs: List[List[int]]) -> Counter:
+        c = Counter()
+        for left, right in zip(shard_seqs, shard_seqs[1:]):
+            if left and right:
+                c[(left[-1], right[0])] += 1
+        return c
+
+    def train(self, corpus: Iterable[str], log_every: int = 128) -> None:
+        """Chunk-atomic BPE with INCREMENTAL pair counting (efficient).
+
+        Pre-tokenization splits each document into atomic chunks (whitespace
+        runs and non-whitespace words, both byte-encoded). Lemma 1 applies at
+        chunk level: C_global = sum over chunks of interior pairs, identical
+        for any shard assignment of chunks. Instead of recounting all pairs
+        at every merge (O(corpus) per merge), an index pair -> {chunk ids}
+        updates only the chunks touched by the merge -- the standard
+        linear-work BPE optimization, with the SAME merge sequence as full
+        recounting (counts are maintained as exact deltas).
+        """
+        import re
+        texts = [t for t in corpus if t]
+        if not texts:
+            raise ValueError("empty corpus")
+        # unique chunks with frequencies (huge speedup; same counts)
+        chunk_pat = re.compile(r"\s+|\S+")
+        freq: Counter = Counter()
+        for t in texts:
+            for chunk in chunk_pat.findall(t):
+                if chunk:
+                    freq[bytes(_byte_encode(chunk))] += 1
+        # de-duplicate identical byte sequences (same bytes, same merges)
+        uniq: Dict[bytes, int] = {}
+        seqs: List[List[int]] = []
+        counts: List[int] = []
+        for seq, c in freq.items():
+            key = bytes(seq)
+            if key in uniq:
+                counts[uniq[key]] += c
+            else:
+                uniq[key] = len(seqs)
+                seqs.append(list(seq))
+                counts.append(c)
+        # shard the UNIQUE chunks (Lemma-1 shard assignment kept explicit)
+        # NOTE: work is single-threaded incremental updates; the shard
+        # assignment documents the equivalence contract (any sharding of the
+        # atomic chunks yields identical merges).
+        assignments: List[List[int]] = [[] for _ in range(self.shards)]
+        for i in range(len(seqs)):
+            assignments[i % self.shards].append(i)
+        _ = assignments
+
+        # global pair counts + pair -> seq ids index
+        total: Counter = Counter()
+        index: Dict[Tuple[int, int], set] = {}
+        for sid, ids in enumerate(seqs):
+            for pair in zip(ids, ids[1:]):
+                total[pair] += counts[sid]
+                index.setdefault(pair, set()).add(sid)
+
+        n_merges = self.vocab_size - 256
+        for step in range(n_merges):
+            if not total:
+                break
+            (a, b), cnt = min(total.items(), key=lambda kv: (-kv[1], kv[0]))
+            if cnt < self.min_freq:
+                break
+            new_id = 256 + step
+            self.vocab[new_id] = self.vocab[a] + self.vocab[b]
+            self.ranks[(a, b)] = step
+            self.merge_left[(a, b)] = a
+            self.merge_right[(a, b)] = b
+            self.merge_history.append((a, b, new_id))
+            # delta-update only the sequences containing the pair
+            affected = index.pop((a, b), set())
+            for sid in affected:
+                ids = seqs[sid]
+                c = counts[sid]
+                # subtract old pair multiset of this sequence
+                for pair in zip(ids, ids[1:]):
+                    total[pair] -= c
+                new_ids = self._apply_merge(ids, a, b, new_id)
+                seqs[sid] = new_ids
+                for pair in zip(new_ids, new_ids[1:]):
+                    total[pair] += c
+                    index.setdefault(pair, set()).add(sid)
+            total.pop((a, b), None)
+            # periodic prune (zero/near-zero entries cannot win the argmax
+            # because min_freq >= 2 > 0, but they would grow memory)
+            if step % 64 == 63:
+                total = Counter({p: v for p, v in total.items() if v > 0})
+            if log_every and (step + 1) % log_every == 0:
+                print(f"[bpe] merge {step+1}/{n_merges} pair=({a},{b}) cnt={cnt}")
+
+    @staticmethod
+    def _count_docs(docs: List[List[int]]) -> Counter:
+        c = Counter()
+        for seq in docs:
+            for pair in zip(seq, seq[1:]):
+                c[pair] += 1
+        return c
+
+    @staticmethod
+    def _apply_merge(seq: List[int], a: int, b: int, new_id: int) -> List[int]:
+        out: List[int] = []
+        i, n = 0, len(seq)
+        while i < n:
+            if i < n - 1 and seq[i] == a and seq[i + 1] == b:
+                out.append(new_id)
+                i += 2
+            else:
+                out.append(seq[i])
+                i += 1
+        return out
+
+    # ---------------------------------------------------------- encode/decode
+    def _encode_seq(self, ids: List[int]) -> List[int]:
+        # greedy by rank (standard BPE encode)
+        while len(ids) >= 2:
+            best, best_rank = None, None
+            for pair in zip(ids, ids[1:]):
+                r = self.ranks.get(pair)
+                if r is not None and (best_rank is None or r < best_rank):
+                    best, best_rank = pair, r
+            if best is None:
+                break
+            a, b = best
+            new_id = 256 + self.ranks[best]
+            ids = self._apply_merge(ids, a, b, new_id)
+        return ids
+
+    def encode(self, text: str) -> List[int]:
+        return self._encode_seq(_byte_encode(text))
+
+    def decode(self, ids: Iterable[int]) -> str:
+        buf = b"".join(self.vocab.get(int(i), bytes([int(i) & 0xFF]))
+                       for i in ids)
+        return buf.decode("utf-8", errors="replace")
+
+    @property
+    def n_vocab(self) -> int:
+        """ACTUAL number of learned ids (256 bytes + merges so far)."""
+        return len(self.vocab)
+
+    # ------------------------------------------------------------------- io
+    def save(self, path: Path) -> None:
+        data = {
+            "vocab_size": self.vocab_size,
+            "min_freq": self.min_freq,
+            "merges": [[a, b, new] for (a, b, new) in self.merge_history],
+        }
+        Path(path).write_text(json.dumps(data), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: Path) -> "ByteBPETokenizer":
+        data = json.loads(Path(path).read_text(encoding="utf-8"))
+        tok = cls(vocab_size=data["vocab_size"], min_freq=data["min_freq"])
+        for a, b, new in data["merges"]:
+            tok.vocab[new] = tok.vocab[a] + tok.vocab[b]
+            tok.ranks[(a, b)] = len(tok.merge_history)
+            tok.merge_left[(a, b)] = a
+            tok.merge_right[(a, b)] = b
+            tok.merge_history.append((a, b, new))
+        return tok