Spaces:
Runtime error
Runtime error
| """ | |
| BPETrainer: learn merge rules from raw text. | |
| Note: this recounts all pairs each iteration (clear but 0(merges x tokens)). | |
| Fine for today's small demo. For tomorrow's full TinyStories run we'll add | |
| incremental pair-count updates so it scales. | |
| """ | |
| import os | |
| import json | |
| from collections import Counter, defaultdict | |
| from .regex_split import split | |
| from .tokenizer import merge_seq | |
| from .special_tokens import SPECIAL_TOKENS, special_token_ids | |
| def _pairs(ids): | |
| """Counter of adjacent id pairs within one word.""" | |
| c = Counter() | |
| for a, b in zip(ids, ids[1:]): | |
| c[(a, b)] += 1 | |
| return c | |
| class BPETrainer: | |
| def __init__(self, vocab_size=16384, verbose=False): | |
| self.vocab_size = vocab_size | |
| self.verbose = verbose | |
| self.merges = [] | |
| self.vocab = {} | |
| # ---- entry points ---- | |
| def train(self, text): | |
| """Train from an in-memory string *used by the test).""" | |
| return self._fit(Counter(split(text))) | |
| def train_from_file(self, path): | |
| """Stream a (possibly multi-GB) text file line by line so we never | |
| hold the whole corpus in memory.""" | |
| word_freqs = Counter() | |
| with open(path, "r", encoding="utf-8") as f: | |
| for i, line in enumerate(f): | |
| if "<|endoftext|>" in line: | |
| line = line.replace("<|endoftext|>", " ") | |
| word_freqs.update(split(line)) | |
| if self.verbose and i and i % 200000 == 0: | |
| print(f" read {i:,} lines, {len(word_freqs):,} unique chunks") | |
| if self.verbose: | |
| print(f" done reading: {len(word_freqs):,} unique chunks") | |
| return self._fit(word_freqs) | |
| # ---- core training ---- | |
| def _fit(self, word_freqs): | |
| num_merges = self.vocab_size - 256 - len(SPECIAL_TOKENS) | |
| if num_merges < 0: | |
| raise ValueError("vocab_size too small for 256 + specials") | |
| # each unique word -> list of byte ids | |
| splits = {w: list(w.encode("utf-8")) for w in word_freqs} | |
| # running pair counts (weighted by word freq) + which words hold each pair | |
| pair_counts = Counter() | |
| where = defaultdict(set) | |
| for w, freq in word_freqs.items(): | |
| for p, c in _pairs(splits[w]).items(): | |
| pair_counts[p] += c * freq | |
| where[p].add(w) | |
| merges = [] | |
| vocab = {i: bytes([i]) for i in range(256)} | |
| for m in range(num_merges): | |
| if not pair_counts: | |
| break | |
| best = max(pair_counts, key=pair_counts.get) | |
| if pair_counts[best] <= 0: | |
| break | |
| new_id = 256 + m | |
| merges.append(best) | |
| vocab[new_id] = vocab[best[0]] + vocab[best[1]] | |
| # only the words containing 'best' can change | |
| for w in list(where[best]): | |
| freq = word_freqs[w] | |
| old = splits[w] | |
| before = _pairs(old) | |
| new = merge_seq(old, best, new_id) | |
| splits[w] = new | |
| after = _pairs(new) | |
| for p in set(before) | set(after): | |
| delta = (after[p] - before[p]) * freq | |
| if delta: | |
| pair_counts[p] += delta | |
| if after[p] > 0: | |
| where[p].add(w) | |
| else: | |
| where[p].discard(w) | |
| if pair_counts.get(p, 0) <= 0: | |
| pair_counts.pop(p, None) | |
| where.pop(p, None) | |
| if self.verbose and (m < 20 or m% 500 == 0): | |
| piece = vocab[new_id].decode("utf-8", errors="replace") | |
| print(f"merge {m:5d}: {best} -> id {new_id} {piece!r}") | |
| self.merges = merges | |
| self.vocab = vocab | |
| return merges, vocab | |
| # ---- save ---- | |
| def save(self, out_dir): | |
| os.makedirs(out_dir, exist_ok=True) | |
| with open(os.path.join(out_dir, "merges.txt"), "w", encoding="utf-8") as f: | |
| for a, b in self.merges: | |
| f.write(f"{a} {b}\n") | |
| # human-readable view (not used for loading, just for inspection) | |
| view = {i: b.decode("utf-8", errors="replace") | |
| for i, b in self.vocab.items()} | |
| for tok, tid in special_token_ids(self.vocab_size).items(): | |
| view[tid] = tok | |
| with open(os.path.join(out_dir, "vocab.json"), "w", encoding="utf-8") as f: | |
| json.dump(view, f, ensure_ascii=False, indent=0) | |
| with open(os.path.join(out_dir, "config.json"), "w", encoding="utf-8") as f: | |
| json.dump({"vocab_size": self.vocab_size}, f) |