Byte Tokenizer and Token-Count Estimator

Study a slow reference byte tokenizer, describe it precisely, then implement a faster `tokenize` method with identical output. The final part estimates token counts under a sampling budget without reimplementing tokenizer logic. Challenge Summary This assessment asks you to build and improve a byte-level tokenizer. This is similar to Byte Pair Encoding (BPE), which is a method used by Large Language Models (LLMs) to process text. The test checks your ability to understand algorithms, improve code performance, and make statistical estimates. There are three parts: Part 0: Read a slow, basic reference solution (slow_tokenize) and explain how it works in plain English. Part 1: Write a faster version (tokenize) that gives the exact same result. Part 2: Write a function to estimate the number of tokens in a text without reading the whole thing. Provided Code Template import random class ByteTokenizer: """ A tokenizer that turns bytes into token IDs. The 'token_alphabet' is a list. The first 256 items are single bytes (0-255). Items after 256 are multi-byte tokens, ordered by when they should merge. """ def __init__(self, token_alphabet: list[bytes]): """ Setup the tokenizer. Args: token_alphabet: A list of byte sequences. Index = token ID. Value = the actual bytes. """ self.token_alphabet = token_alphabet self.token_to_id: dict[bytes, int] = { token: i for i, token in enumerate(token_alphabet) } # TODO: Add any setup logic here for Part 1 and Part 2 # NOTE: DO NOT CHANGE ANY CODE IN slow_tokenize. def slow_tokenize(self, text: bytes) -> list[int]: """ Tokenize text using a slow, simple method. Returns a list of integers (token IDs). """ # STEP 1: INITIALIZATION # # Start by treating every byte in the text as a separate token. # Since the first 256 tokens match the byte values 0-255, # we can just turn the text into a list of integers. # # Example: b"hello" becomes [104, 101, 108, 108, 111]. token_ids: list[int] = list(text) # STEP 2: MERGING INTO MULTI-BYTE TOKENS # # Go through the multi-byte tokens one by one, starting from ID 256. # For each token, scan the current list and see if we can merge # two smaller tokens into this bigger one. for cur_token_id in range(256, len(self.token_alphabet)): cur_token_bytes = self.token_alphabet[cur_token_id] # Find pairs of smaller tokens that make up this current token. pairs_to_merge: set[tuple[int, int]] = set() for split_point in range(1, len(cur_token_bytes)): left_bytes = cur_token_bytes[:split_point] right_bytes = cur_token_bytes[split_point:] left_id = self.token_to_id.get(left_bytes) right_id = self.token_to_id.get(right_bytes) # If both parts exist as tokens, we can merge them. if left_id is not None and right_id is not None: pairs_to_merge.add((left_id, right_id)) # Do the merges from left to right. new_token_ids: list[int] = [] i = 0 while i < len(token_ids): # Check if the pair at (i, i+1) matches a valid merge pair. if (i < len(token_ids) - 1 and (token_ids[i], token_ids[i + 1]) in pairs_to_merge): new_token_ids.append(cur_token_id) i += 2 # Skip the next item because it was merged else: # Cannot merge, keep the original token new_token_ids.append(token_ids[i]) i += 1 token_ids = new_token_ids return token_ids def tokenize(self, text: bytes) -> list[int]: """ Tokenize text efficiently. Must give the same result as slow_tokenize. """ # TODO: Implement this in Part 1 raise NotImplementedError def estimate_token_count( self, text: bytes, sample_size: int, rng: random.Random, ) -> int: """ Guess the number of tokens without processing the whole text. Constraints: - Total bytes passed to tokenize() must be <= sample_size. - sample_size is always at least 1000. - If the text is small (<= sample_size), return the exact count. - Use rng for random numbers. """ # TODO: Implement this in Part 2 raise NotImplementedError Part 0: Algorithm Description Your Goal Read the slow_tokenize function. Write a simple description of how it works. A student should be able to write the code just by reading your description. How the Algorithm Works The algorithm turns bytes into a list of token IDs in two main steps: Step 1 — Setup: First, turn the text into a list of numbers. Every byte becomes its own token. Since the first 256 tokens in our list match the byte values (0 to 255), this is easy. For example, the text b"hello" becomes [104, 101, 108, 108, 111]. Step 2 — Merging Loop: We look at every multi-byte token in our dictionary, starting from ID 256 and going up. For each token, we do the following: Find Valid Pairs: We check how this token can be built. We split the token's bytes into two parts (left and right). If both parts are already known tokens, we save that pair of IDs as a "merge rule." Scan and Replace: We look through our current list of numbers from left to right. If we see two numbers next to each other that match one of our merge rules: We replace those two numbers with the new, bigger token ID. We skip ahead past the two numbers we just used. If they don't match, we keep the number as is and move one step forward. Repeat: We update our list and move to the next token ID in the dictionary. Important Details: One big token might be made from different pairs (e.g., "ABC" could be "A"+"BC" or "AB"+"C"). We always merge from left to right. If we have A, B, A, B and merge (A, B), we get two new tokens. If we have A, B, B, we only merge the first pair. Order matters. We must process token 256 before token 257. Part 1: Fast Tokenize Your Goal Write the tokenize method. It must give the exact same output as slow_tokenize but run much faster. Target: Process ~500KB of text with 10,000+ tokens in under 10 seconds. Performance Bottlenecks The original code is slow because it has a time complexity of O(V × N). V is the number of tokens (vocabulary size). N is the length of the text. It scans the entire text for every single token in the vocabulary. Even if a token never appears in the text, the code still scans for it. This wastes time. The Optimization Strategy We can skip the useless scans using two data structures: Min-Heap (Priority Queue): Instead of checking tokens in order (256, 257...), we let a heap tell us which merge needs to happen next. We store all possible adjacent pairs in the heap. The "priority" is the target token ID (lower IDs merge first). Doubly-Linked List: Merging items in a standard Python list is slow because you have to shift all the other items. A linked list lets us remove two items and insert a new one in O(1) time. This changes the complexity to O(N log N). We only do work when a merge actually happens. Why This Strategy Works Using a heap ensures we follow the rules: Correct Order: The heap always gives us the merge with the lowest target ID first. This matches the original loop for cur_token_id in range(256, ...). Left-to-Right: If we have multiple merges for the same ID, we process the leftmost one first based on its position index. Optimized Code Solution import heapq import random class ByteTokenizer: def __init__(self, token_alphabet: list[bytes]): self.token_alphabet = token_alphabet self.token_to_id: dict[bytes, int] = { token: i for i, token in enumerate(token_alphabet) } # Pre-calculation: # Find out which pair of IDs creates which new token ID. self.merge_rules: dict[tuple[int, int], int] = {} for token_id in range(256, len(self.token_alphabet)): token_bytes = self.token_alphabet[token_id] for split_point in range(1, len(token_bytes)): left_bytes = token_bytes[:split_point] right_bytes = token_bytes[split_point:] left_id = self.token_to_id.get(left_bytes) right_id = self.token_to_id.get(right_bytes) if left_id is not None and right_id is not None: pair = (left_id, right_id) # If a pair maps to multiple tokens (rare), # we want the one with the lowest ID. if pair not in self.merge_rules: self.merge_rules[pair] = token_id else: self.merge_rules[pair] = min( self.merge_rules[pair], token_id ) def tokenize(self, text: bytes) -> list[int]: if not text: return [] n = len(text) token_ids = list(text) # Build a doubly-linked list using arrays. # prev_node[i] stores the index of the item before i. # next_node[i] stores the index of the item after i. prev_node = list(range(-1, n - 1)) # -1 means no previous node next_node = list(range(1, n + 1)) # n means end of list alive = [True] * n # Tracks if a node has been merged/deleted # Add all adjacent pairs to the heap. heap: list[tuple[int, int]] = [] for i in range(n - 1): pair = (token_ids[i], token_ids[i + 1]) if pair in self.merge_rules: # Store (target_token_id, position) heapq.heappush(heap, (self.merge_rules[pair], i)) while heap: target_id, pos = heapq.heappop(heap) # Check if this merge is still valid. if not alive[pos]: continue nxt = next_node[pos] if nxt >= n or not alive[nxt]: continue # Check if the pair still matches the target. pair = (token_ids[pos], token_ids[nxt]) if self.merge_rules.get(pair) != target_id: continue # Do the merge: # Update the token at 'pos' and mark 'nxt' as dead. token_ids[pos] = target_id alive[nxt] = False # Update linked list pointers to skip 'nxt'. next_node[pos] = next_node[nxt] if next_node[nxt] < n: prev_node[next_node[nxt]] = pos # Check if the new token creates a merge with its LEFT neighbor. p = prev_node[pos] if p >= 0: new_pair = (token_ids[p], token_ids[pos]) if new_pair in self.merge_rules: heapq.heappush(heap, (self.merge_rules[new_pair], p)) # Check if the new token creates a merge with its RIGHT neighbor. nxt2 = next_node[pos] if nxt2 < n: new_pair = (token_ids[pos], token_ids[nxt2]) if new_pair in self.merge_rules: heapq.heappush(heap, (self.merge_rules[new_pair], pos)) # Collect the final result by following the linked list. result: list[int] = [] i = 0 while i < n: result.append(token_ids[i]) i = next_node[i] return result Step-by-Step Example Imagine tokenizing b"aabaa". Tokens 0..255: Single bytes. Token 256: b"aa" (made of a + a). Token 257: b"aab" (made of aa + b). Start: List: [a, a, b, a, a] (indices 0, 1, 2, 3, 4) Heap finds pairs (a,a) at index 0 and index 3. Both target token 256. Step 1: Pop (256, 0) from heap. Merge index 0 and 1. List becomes: [aa, b, a, a] (Indices used: 0, 2, 3, 4. Index 1 is dead). New check: Does aa (index 0) + b (index 2) merge? Yes! It makes 257. Add to heap. Step 2: Pop (256, 3) from heap. Merge index 3 and 4. List becomes: [aa, b, aa] (Indices used: 0, 2, 3. Index 4 is dead). Step 3: Pop (257, 0) from heap. Merge index 0 (aa) and 2 (b). List becomes: [aab, aa] (Indices used: 0, 3. Index 2 is dead). Result: [257, 256]. This is correct. Time and Space Complexity Metric Slow Version Fast Version Time O(V × N) O(N log N) Space O(N) O(N) Pre-calc None O(V × L) (L = token length) The fast version is much better when the vocabulary (V) is large. Part 2: Estimate Token Count Your Goal Write estimate_token_count. You need to guess how many tokens the full text will have, but you can only process a small amount of data (sample_size). Accuracy: Within 20% generally, within 5% for larger samples. Sampling Strategy We use Stratified Sampling. This means we don't just pick random spots. We divide the text into equal sections and take a sample from each section. Divide text into segments. Take a small "chunk" from each segment. Tokenize just those chunks. Calculate the Tokens-Per-Byte Ratio. Multiply that ratio by the total file size. Why not random sampling? Random sampling might pick all chunks from the beginning of the file (which might be headers or imports). Stratified sampling ensures we see code, comments, and data from all over the file. Estimation Code Solution def estimate_token_count( self, text: bytes, sample_size: int, rng: random.Random, ) -> int: # If the text is small, just count it exactly. if sample_size >= len(text): return len(self.tokenize(text)) # Decide how big each chunk should be and how many to take. # We want at least 4 chunks for good variety. # Chunks should be at least 250 bytes. chunk_size = max(250, min(sample_size // 4, len(text) // 10)) num_samples = sample_size // chunk_size # Safety check if num_samples < 1: num_samples = 1 chunk_size = sample_size segment_size = len(text) // num_samples total_tokens = 0 total_bytes = 0 for i in range(num_samples): # Pick a random starting point within the i-th segment seg_start = i * segment_size seg_end = min(seg_start + segment_size, len(text)) # Ensure we have room for a full chunk max_start = max(seg_start, seg_end - chunk_size) start = rng.randint(seg_start, max_start) end = min(start + chunk_size, len(text)) chunk = text[start:end] tokens = self.tokenize(chunk) total_tokens += len(tokens) total_bytes += len(chunk) # Calculate average and extrapolate ratio = total_tokens / total_bytes return round(ratio * len(text)) Logic Behind the Estimation The number of tokens per byte is usually consistent in a document. For example, Python code usually compresses at a similar rate throughout a file. There is a small error at the "cut points" where we slice the chunk (we might cut a token in half), but since the chunks are reasonably large (250+ bytes) and tokens are small (3-4 bytes), this error is less than 1%. Step-by-Step Example Imagine a 100,000 byte file. sample_size is 10,000. Calculate Chunk Size: We can afford 4 chunks of 2,500 bytes. Divide: We split the file into 4 zones (0-25k, 25k-50k, etc.). Sample: We pick a random 2,500 byte chunk from each zone. Count: We tokenize those chunks. Let's say we find 3,500 tokens total in our 10,000 bytes. Ratio: 3,500 / 10,000 = 0.35 tokens per byte. Final Math: 0.35 * 100,000 (total size) = 35,000 tokens. Important Takeaways Understanding Merge Order The order matters. BPE learns merges based on frequency. Token 256 is the most common pair, 257 is the second most common. Merging in this exact order (256, then 257...) guarantees we get the same result as the training phase. Other Possible Solutions Trie (Prefix Tree): You could use a Trie to find the longest match. However, BPE doesn't always want the longest match; it wants the highest priority match. Segment Tree: Good for data that changes often, but too complex for this problem since the text is static. Making Estimates Better To improve Part 2, you could: Overlap Correction: Sample slightly overlapping chunks to see how tokens behave at the edges. Weighted Sampling: If you know parts of the file are weird (like binary data), sample them more. Real-World Applications Tools like tiktoken (used by OpenAI) or SentencePiece use this same logic but written in C++ or Rust for maximum speed. They also use HashMaps to look up pairs instantly, just like our merge_rules dictionary.