Tokenizer Design Choices: BPE, Vocabulary, and Implementation

# Start with bytes (256 base tokens)
text = "hello hello world"
tokens = [104, 101, 108, 108, 111, 32, ...]  # byte values
 
# Find most frequent pair
pairs = count_pairs(tokens)  # (104, 101) appears most
# => merge (h, e) -> token 256
 
# Repeat for vocab_size - 256 merges
# Final vocabulary: [0..255, "he", "ll", "lo", "hello", ...]

def train_bpe(text_chunks, vocab_size):
    """Core BPE training loop"""
    # 1. Initialize with byte-level tokens
    words = [chunk.encode("utf-8") for chunk in text_chunks]
    vocab = list(range(256))  # base vocabulary
    merges = {}  # (token_a, token_b) -> merged_token_id
    
    # 2. Iteratively merge most frequent pairs
    for i in range(vocab_size - 256):
        # Count all adjacent pairs
        pair_counts = count_pairs_in_corpus(words)
        
        # Find most frequent pair
        best_pair = max(pair_counts, key=pair_counts.get)
        
        # Create new token for this pair
        new_token_id = 256 + i
        merges[best_pair] = new_token_id
        vocab.append(best_pair)
        
        # Apply merge to all words
        words = apply_merge(words, best_pair, new_token_id)
    
    return vocab, merges

SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""

# nanochat uses \p{N}{1,2} (1-2 digits)
# GPT-4 uses \p{N}{1,3} (1-3 digits)

tokenizer = HFTokenizer(BPE(
    byte_fallback=True,  # Enable byte-level fallback
    unk_token=None,      # No unknown token needed!
    fuse_unk=False,
))

def text_iterator():
    """
    Stream text from parquet files:
    1) Flatten batches into single iterator
    2) Crop documents to doc_cap characters
    3) Break after max_chars total
    """
    nchars = 0
    for batch in parquets_iter_batched(split="train"):
        for doc in batch:
            # Crop long documents
            doc_text = doc[:args.doc_cap] if len(doc) > args.doc_cap else doc
            nchars += len(doc_text)
            yield doc_text
            if nchars > args.max_chars:
                return

python scripts/tok_train.py \
    --vocab_size 65536 \
    --max_chars 10_000_000_000 \
    --doc_cap 10000

token_bytes = []
for token_id in range(vocab_size):
    token_str = tokenizer.decode([token_id])
    if token_str in special_tokens:
        token_bytes.append(0)  # Special tokens don't count
    else:
        token_bytes.append(len(token_str.encode("utf-8")))
 
# Save for bits-per-byte evaluation
torch.save(torch.tensor(token_bytes), "token_bytes.pt")

# Traditional loss: depends on vocabulary size
loss = -log P(token_id)  # varies with vocab_size
 
# Bits-per-byte: normalized by token byte count
bpb = loss / token_bytes[token_id]  # comparable across tokenizers

struct MergeJob {
    pair: (u32, u32),           // token pair
    count: u64,                 // frequency
    pos: AHashSet<usize>,       // word indices where pair occurs
}
 
fn train_core_incremental(&mut self, words: Vec<Word>, counts: Vec<i32>, vocab_size: u32) {
    // 1. Count all pairs in parallel
    let (pair_counts, where_to_update) = count_pairs_parallel(&words, &counts);
    
    // 2. Build max-heap of pairs by frequency
    let mut heap = OctonaryHeap::with_capacity(pair_counts.len());
    for (pair, pos) in where_to_update {
        heap.push(MergeJob { pair, count: pair_counts[pair], pos });
    }
    
    // 3. Merge loop: process vocab_size - 256 pairs
    for i in 0..(vocab_size - 256) {
        // Pop most frequent pair
        let top = heap.pop().unwrap();
        
        // Lazy refresh: check if count is still valid
        if top.count != pair_counts[top.pair] {
            heap.push(MergeJob { count: pair_counts[top.pair], ..top });
            continue;
        }
        
        // Record merge
        let new_id = 256 + i;
        self.merges.insert(top.pair, new_id);
        
        // Apply merge to affected words
        for &word_idx in &top.pos {
            let deltas = words[word_idx].merge_pair(top.pair, new_id);
            
            // Update pair counts based on deltas
            for (pair, delta) in deltas {
                pair_counts[pair] += delta * counts[word_idx];
                if delta > 0 {
                    heap.push(MergeJob { pair, count: pair_counts[pair], .. });
                }
            }
        }
    }
}

use dary_heap::OctonaryHeap;  // 8-ary heap vs binary

// Pop from heap
let top = heap.pop();
 
// Check if count is stale (other merges updated it)
if top.count != pair_counts[top.pair] {
    // Re-push with updated count (lazy refresh)
    heap.push(MergeJob { count: pair_counts[top.pair], ..top });
    continue;
}

fn count_pairs_parallel(words: &[Word], counts: &[i32]) -> (HashMap<Pair, i32>, HashMap<Pair, HashSet<usize>>) {
    words.par_iter()  // Rayon parallel iterator
        .enumerate()
        .map(|(i, word)| {
            // Count pairs in this word
            let mut local_counts = HashMap::new();
            for (a, b) in word.pairs() {
                *local_counts.entry((a, b)).or_default() += counts[i];
            }
            local_counts
        })
        .reduce(/* merge local counts */)
}

fn merge_pair(&mut self, pair: (a, b), new_id) -> Vec<(Pair, i32)> {
    let mut deltas = Vec::new();
    
    // For each occurrence of (a, b):
    // - Remove pairs: (left, a), (a, b), (b, right)
    // - Add pairs: (left, new_id), (new_id, right)
    
    if let Some(left) = left_neighbor {
        deltas.push(((left, a), -1));      // removed
        deltas.push(((left, new_id), +1)); // added
    }
    deltas.push(((a, b), -1));             // removed
    if let Some(right) = right_neighbor {
        deltas.push(((b, right), -1));     // removed
        deltas.push(((new_id, right), +1)); // added
    }
    
    deltas
}

pub fn train_from_iterator(&mut self, py: Python, iterator: &PyAny, vocab_size: u32, buffer_size: usize) {
    let mut buf: Vec<String> = Vec::with_capacity(buffer_size);
    
    loop {
        // 1. Refill buffer (under GIL)
        let exhausted = refill(&mut buf, iterator)?;
        
        // 2. Process buffer (release GIL, parallel)
        let local_counts = py.allow_threads(|| {
            buf.par_iter()
                .map(|text| {
                    // Apply regex, count chunks
                    let mut counts = HashMap::new();
                    for chunk in pattern.find_iter(text) {
                        *counts.entry(chunk).or_default() += 1;
                    }
                    counts
                })
                .reduce(/* merge */)
        });
        
        // 3. Merge into global counts
        for (chunk, count) in local_counts {
            *global_counts.entry(chunk).or_default() += count;
        }
        
        if exhausted { break; }
    }
}

# After training with RustBPE
pattern = tokenizer.get_pattern()
mergeable_ranks_list = tokenizer.get_mergeable_ranks()
mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list}
 
# Add special tokens
special_tokens = {
    "<|bos|>": 65536,
    "<|user_start|>": 65537,
    "<|user_end|>": 65538,
    # ... 9 special tokens total
}
 
# Create tiktoken Encoding
enc = tiktoken.Encoding(
    name="rustbpe",
    pat_str=pattern,
    mergeable_ranks=mergeable_ranks,
    special_tokens=special_tokens,
)

# Single string
tokens = enc.encode_ordinary("Hello, world!")
 
# Batch encoding (parallel)
texts = ["text1", "text2", ..., "text1000"]
tokens_batch = enc.encode_ordinary_batch(texts, num_threads=8)

# Ordinary encoding: treats special tokens as text
tokens = enc.encode_ordinary("<|bos|>Hello")
# => [60, 124, 98, 111, 115, 124, 62, 9906]  (raw bytes)
 
# Special encoding: recognizes special token
tokens = enc.encode("<|bos|>Hello", allowed_special="all")
# => [65536, 9906]  (special token ID + "Hello")
 
# Single special token
bos_id = enc.encode_single_token("<|bos|>")  # => 65536

SPECIAL_TOKENS = [
    "<|bos|>",           # Beginning of sequence (every document)
    "<|user_start|>",    # User message delimiter
    "<|user_end|>",
    "<|assistant_start|>",  # Assistant message delimiter
    "<|assistant_end|>",
    "<|python_start|>",  # Tool use: Python REPL
    "<|python_end|>",
    "<|output_start|>",  # Tool output
    "<|output_end|>",
]

conversation = {
    "messages": [
        {"role": "user", "content": "What is 2+2?"},
        {"role": "assistant", "content": "2+2 equals 4."},
    ]
}
 
# Need to produce:
# tokens: [<bos>, <user_start>, "What", "is", "2", "+", "2", "?", <user_end>,
#          <assistant_start>, "2", "+", "2", "equals", "4", ".", <assistant_end>]
# mask:   [0,     0,            0,     0,   0,  0,  0,  0,  0,
#          0,                  1,   1,   1,   1,       1,   1,   1]

def render_conversation(self, conversation, max_tokens=2048):
    """
    Tokenize a conversation and return (ids, mask).
    - ids: list of token IDs
    - mask: 1 for assistant tokens (supervised), 0 otherwise
    """
    ids, mask = [], []
    
    def add_tokens(token_ids, mask_val):
        if isinstance(token_ids, int):
            token_ids = [token_ids]
        ids.extend(token_ids)
        mask.extend([mask_val] * len(token_ids))
    
    # Fetch special token IDs
    bos = self.encode_special("<|bos|>")
    user_start = self.encode_special("<|user_start|>")
    user_end = self.encode_special("<|user_end|>")
    assistant_start = self.encode_special("<|assistant_start|>")
    assistant_end = self.encode_special("<|assistant_end|>")
    
    # Add BOS (unsupervised)
    add_tokens(bos, 0)
    
    # Process messages
    for i, message in enumerate(conversation["messages"]):
        role = message["role"]
        content = message["content"]
        
        if role == "user":
            add_tokens(user_start, 0)
            add_tokens(self.encode(content), 0)
            add_tokens(user_end, 0)
            
        elif role == "assistant":
            add_tokens(assistant_start, 0)  # Start token not supervised
            add_tokens(self.encode(content), 1)  # Content IS supervised
            add_tokens(assistant_end, 1)  # End token IS supervised
    
    # Truncate to max_tokens
    ids = ids[:max_tokens]
    mask = mask[:max_tokens]
    
    return ids, mask

{
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "Hello!"},
        {"role": "assistant", "content": "Hi there!"},
    ]
}

if conversation["messages"][0]["role"] == "system":
    conversation = copy.deepcopy(conversation)
    messages = conversation["messages"]
    assert messages[1]["role"] == "user"
    # Prepend system content to user message
    messages[1]["content"] = messages[0]["content"] + "\n\n" + messages[1]["content"]
    messages = messages[1:]  # Remove system message

{
    "role": "assistant",
    "content": [
        {"type": "text", "text": "Let me calculate that:"},
        {"type": "python", "text": "print(2+2)"},
        {"type": "python_output", "text": "4"},
        {"type": "text", "text": "The answer is 4."},
    ]
}

add_tokens(assistant_start, 0)
 
for part in content:
    value_ids = self.encode(part["text"])
    
    if part["type"] == "text":
        add_tokens(value_ids, 1)  # Supervised
        
    elif part["type"] == "python":
        add_tokens(python_start, 1)
        add_tokens(value_ids, 1)  # Supervised: learn to generate code
        add_tokens(python_end, 1)
        
    elif part["type"] == "python_output":
        add_tokens(output_start, 0)  # NOT supervised
        add_tokens(value_ids, 0)     # Output comes from Python, not model
        add_tokens(output_end, 0)
 
add_tokens(assistant_end, 1)

def visualize_tokenization(self, ids, mask):
    """Color-code tokens by supervision mask"""
    RED = '\033[91m'    # Unsupervised
    GREEN = '\033[92m'  # Supervised
    RESET = '\033[0m'
    
    tokens = []
    for token_id, mask_val in zip(ids, mask):
        token_str = self.decode([token_id])
        color = GREEN if mask_val == 1 else RED
        tokens.append(f"{color}{token_str}{RESET}")
    
    return '|'.join(tokens)

text = "Hello, world!"
encoded = tokenizer.encode(text)
encoded_bytes = text.encode('utf-8')
 
compression_ratio = len(encoded_bytes) / len(encoded)
# Higher ratio = fewer tokens for same text = better compression

# Load tokenizers
gpt2_tok = RustBPETokenizer.from_pretrained("gpt2")
gpt4_tok = RustBPETokenizer.from_pretrained("cl100k_base")
ours_tok = get_tokenizer()
 
# Test on diverse data
test_data = [
    ("news", news_article),
    ("korean", korean_text),
    ("code", python_code),
    ("math", latex_math),
    ("science", technical_prose),
]
 
# Encode and compare
for name, text in test_data:
    gpt2_tokens = gpt2_tok.encode(text)
    gpt4_tokens = gpt4_tok.encode(text)
    ours_tokens = ours_tok.encode(text)
    
    gpt2_ratio = len(text.encode('utf-8')) / len(gpt2_tokens)
    gpt4_ratio = len(text.encode('utf-8')) / len(gpt4_tokens)
    ours_ratio = len(text.encode('utf-8')) / len(ours_tokens)
    
    print(f"{name:10} GPT-2: {gpt2_ratio:.2f}  GPT-4: {gpt4_ratio:.2f}  Ours: {ours_ratio:.2f}")

# Load token_bytes mapping
token_bytes = get_token_bytes(device="cuda")  # [vocab_size]
 
# During evaluation
losses = F.cross_entropy(logits, targets, reduction='none')  # [batch, seq_len]
 
# Weight each loss by token byte count
token_bytes_flat = token_bytes[targets.flatten()]  # [batch*seq_len]
valid_mask = token_bytes_flat > 0  # Exclude special tokens
 
# Bits-per-byte
bpb = (losses.flatten()[valid_mask] / token_bytes_flat[valid_mask]).mean()

class HuggingFaceTokenizer:
    """Pure Python, uses HF tokenizers library"""
    
    @classmethod
    def train_from_iterator(cls, text_iterator, vocab_size):
        tokenizer = HFTokenizer(BPE(byte_fallback=True))
        # ... configure pre-tokenizer, decoder
        trainer = BpeTrainer(vocab_size=vocab_size, special_tokens=SPECIAL_TOKENS)
        tokenizer.train_from_iterator(text_iterator, trainer)
        return cls(tokenizer)

class RustBPETokenizer:
    """Rust training, tiktoken inference"""
    
    @classmethod
    def train_from_iterator(cls, text_iterator, vocab_size):
        # Train with Rust
        tokenizer = rustbpe.Tokenizer()
        tokenizer.train_from_iterator(text_iterator, vocab_size, pattern=SPLIT_PATTERN)
        
        # Export to tiktoken
        pattern = tokenizer.get_pattern()
        mergeable_ranks = tokenizer.get_mergeable_ranks()
        enc = tiktoken.Encoding(name="rustbpe", pat_str=pattern, mergeable_ranks=mergeable_ranks, ...)
        return cls(enc)

# Good: Powers of 2 for efficiency
vocab_sizes = [32768, 65536, 131072]
 
# Avoid: Arbitrary sizes
vocab_sizes = [50000, 75000]  # Worse for hardware alignment

# Good: Cap individual documents
def text_iterator():
    for doc in dataset:
        yield doc[:10000]  # Prevent single huge doc from dominating
 
# Bad: No capping
def text_iterator():
    for doc in dataset:
        yield doc  # 100MB document consumes 1% of training data alone

# Good: Explicit, unambiguous delimiters
SPECIAL_TOKENS = ["<|bos|>", "<|user_start|>", "<|user_end|>"]
 
# Bad: Ambiguous markers (could appear in text)
SPECIAL_TOKENS = ["[BOS]", "[USER]", "[/USER]"]

# Good: Only supervise assistant tokens
if role == "assistant":
    add_tokens(content_ids, mask=1)  # Supervised
else:
    add_tokens(content_ids, mask=0)  # Not supervised
 
# Bad: Supervise everything (including user messages)
add_tokens(content_ids, mask=1)  # Model learns to generate user messages

# Bad: No byte fallback
tokenizer = BPE(byte_fallback=False, unk_token="<UNK>")
# Problem: Unknown characters → <UNK> token → information loss
 
# Good: Byte fallback enabled
tokenizer = BPE(byte_fallback=True, unk_token=None)
# Solution: Unknown patterns → individual bytes (no information loss)

# Bad: Different patterns for training vs inference
train_pattern = r"\w+|\d+|[^\w\s]+"
inference_pattern = r"\w+|\s+|."  # Oops, different!
 
# Good: Store pattern with tokenizer
tokenizer.pattern = SPLIT_PATTERN  # Use same pattern always

# Bad: Treating special tokens as ordinary text
text = "<|bos|>Hello"
tokens = tokenizer.encode(text)  # Encodes "<|bos|>" as 5-6 tokens
 
# Good: Explicit special token injection
tokens = [tokenizer.encode_special("<|bos|>")] + tokenizer.encode("Hello")
# => [65536, 9906]  (correct)

# Bad: Using raw loss for evaluation
loss = F.cross_entropy(logits, targets).mean()
# Problem: Loss depends on vocab_size (not comparable)
 
# Good: Normalize by token bytes
token_bytes_flat = token_bytes[targets]
bpb = (loss / token_bytes_flat).mean()
# Solution: Vocabulary-agnostic metric

# Bad: Truncate but don't track actual lengths
ids = ids[:max_tokens]  # Lost information about original length
 
# Good: Track lengths or use attention masks
ids = ids[:max_tokens]
actual_length = min(len(original_ids), max_tokens)
# Use actual_length for loss masking

# 1. Visualize tokenization
viz = tokenizer.visualize_tokenization(ids, mask)
print(viz)  # Color-coded output
 
# 2. Decode individual tokens
for token_id in ids:
    print(f"{token_id}: {tokenizer.decode([token_id])!r}")
 
# 3. Check special token IDs
special_ids = {name: tokenizer.encode_special(name) for name in SPECIAL_TOKENS}
print(special_ids)
 
# 4. Verify encode/decode round-trip
text = "Test 你好 123"
tokens = tokenizer.encode(text)
decoded = tokenizer.decode(tokens)
assert decoded == text, f"Round-trip failed: {text!r} != {decoded!r}"