LLM Developer DocsLLM Docs

Loading the search index ...

↑ ↓ move, Enter open, Esc close

GitHub

tokenizer.py

Byte-pair-encoding (BPE) tokenizer, working on characters

Where it fits

10 of 10 functions, classes and methods here are reachable from train.main, and 6 from talk.main.

Pipeline stageStepExplained on
Text to training windowsTokenizertokenization

Imports and used by

Imports from this repository
None
Other imports
re
Imported by
checkpoint.py, data.py, tests/test_checkpoint.py, tests/test_tokenizer.py, train.py
Used by, in the pipeline
checkpoint.py, data.py, talk.py, train.py
Used by, in tests
tests/test_checkpoint.py, tests/test_tokenizer.py

Mathematics used

ConceptWhat this file doesLinesStatus
Pre-tokens and word countsComputes it8, 9Implemented
Applying a mergeComputes it, Calls it16, 17, 57, 70Implemented
Pair frequency and the best pairComputes it35, 37, 38, 46, 48Implemented
Encode and decodeComputes it, Calls it70, 89, 95, 96Implemented
The vocabularyComputes it102, 105, 110Implemented

Functions and constants

def encode(self, text)Lines 86-90GitHub

Text to ids; tokens not in the vocabulary map to the <unk> id

Reached from
train.main, talk.main
Used by, in the pipeline
data.encode_corpus, talk.main
Used by, in tests
tests.test_checkpoint.test_loaded_model_gives_identical_logits, tests.test_checkpoint.test_non_ascii_characters_survive, tests.test_checkpoint.test_tokenizer_round_trips, tests.test_tokenizer.test_no_characters_dropped, tests.test_tokenizer.test_round_trip, tests.test_tokenizer.test_unseen_character_becomes_unk
Calls
tokenizer.apply_merges
Mathematics
Encode and decode
ConstantLines
UNK = "<unk>"3
PRE_TOKEN_RE = re.compile(r"\w+|\s|[^\w\s]")4

Source

tokenizer.py119 linesGitHub
import reUNK = "<unk>"PRE_TOKEN_RE = re.compile(r"\w+|\s|[^\w\s]")def count_words(text):    counts = {}    for word in PRE_TOKEN_RE.findall(text):        counts[word] = counts.get(word, 0) + 1    return countsdef merge_tokens(tokens, pair, merged):    result = []    i = 0    while i < len(tokens):        if i < len(tokens) - 1 and tokens[i] == pair[0] and tokens[i + 1] == pair[1]:            result.append(merged)            i += 2        else:            result.append(tokens[i])            i += 1    return resultdef train_bpe(word_freqs, max_merges=1000):    word_splits = {}    for word in word_freqs:        word_splits[word] = list(word)    merges = []    for step in range(max_merges):        pair_freq = {}        for word in word_splits:            tokens = word_splits[word]            weight = word_freqs[word]            for i in range(len(tokens) - 1):                key = (tokens[i], tokens[i + 1])                pair_freq[key] = pair_freq.get(key, 0) + weight        if len(pair_freq) == 0:            break        best_pair = None        best_count = 0        for pair in pair_freq:            if pair_freq[pair] > best_count:                best_count = pair_freq[pair]                best_pair = pair        if best_pair is None:            break        a, b = best_pair        merged = a + b        for word in word_splits:            word_splits[word] = merge_tokens(word_splits[word], best_pair, merged)        merges.append({"pair": best_pair, "merged": merged, "frequency": best_count})    return mergesdef apply_merges(text, merges):    result = []    for pre_token in PRE_TOKEN_RE.findall(text):        tokens = list(pre_token)        for m in merges:            tokens = merge_tokens(tokens, m["pair"], m["merged"])        result.extend(tokens)    return resultdef train_bpe_on_text(text, max_merges=1000):    return train_bpe(count_words(text), max_merges)class Tokenizer:    def __init__(self, merges, stoi, itos):        self.merges = merges        self.stoi = stoi        self.itos = itos        self.vocab_size = len(itos)    def encode(self, text):        ids = []        for token in apply_merges(text, self.merges):            ids.append(self.stoi.get(token, self.stoi[UNK]))        return ids    def decode(self, ids):        tokens = []        for i in ids:            tokens.append(self.itos[i])        return "".join(tokens)def build_tokenizer(text, max_merges=1000):    merges = train_bpe_on_text(text, max_merges)    itos = [UNK]    seen = {UNK}    for ch in text:        if ch not in seen:            seen.add(ch)            itos.append(ch)    for m in merges:        if m["merged"] not in seen:            seen.add(m["merged"])            itos.append(m["merged"])        stoi = {}    for i in range(len(itos)):        stoi[itos[i]] = i    return Tokenizer(merges, stoi, itos)