LLM Loves Tokenizers! Implementing BPE from Zero

LLM is the main tool behind most of the things nowadays. For example, an agent has different loops connected via graphs, but the brain is LLM. When we want to do something, we always go for LLM. In that , we have a lot of stages where the tokenizer is the main and first thing.

Tokenizers are the one which converts raw strings into numbers, so that the machine can process further by converting them into encodings to understand the context and to work on it efficiently, to do the math with them.

Tiktoken and SentencePiece are the tokenizers behind modern LLMs. Google developed and uses SentencePiece. And tiktoken is used by the GPT series, byte pair encoding is the main algorithm behind tiktoken.

Byte pair encoding is an algorithm that maps unicode to characters. When a text is given to the LLM, it splits the text into chunks, and for each chunk, a tokenID is generated. It then takes the list of IDs for a word, finds the most frequently occurring pair, and assigns a new token to it. By continuously doing this, a word like “hello” can be squished into an array of with one or two token IDs. In the end, the tokenizer trains on a vocabulary of a fixed size and remembers it. The vocabulary size determines the number of tokens the tokenizer knows. For example, GPT-4 has a vocabulary size of around 100,000.

THE CODE

def __init__(self, vocab_size=32000):
        self.vocab_size = vocab_size
        self.merges = {}
        self.vocab = {}
        self.id_to_token = {}

this the intialization part where we define the vocab size and vocab id_to_token.

 def _bytes_to_unicode(self):
        bs = list(range(ord("!"), ord("~")+1)) + list(range(ord("¡"), ord("¬")+1)) + list(range(ord("®"), ord("ÿ")+1))
        cs = bs[:]
        n = 0
        for b in range(2**8):
            if b not in bs:
                bs.append(b)
                cs.append(2**8 + n)
                n += 1
        return {b: chr(c) for b, c in zip(bs, cs)}

this the basic assigning of charcters to its unicode where whe fisrt assign the basic chartes and their unocde.and for the oothers we assign to a newer unqiye character.

def _get_stats(self, tokens):
        pairs = Counter()
        for token_seq in tokens:
            for i in range(len(token_seq) - 1):
                pair = (token_seq[i], token_seq[i + 1])
                pairs[pair] += 1
        return pairs

    def _merge_pair(self, a, b, new_id, tokens):
        new_tokens = []
        for token_seq in tokens:
            new_seq = []
            i = 0
            while i < len(token_seq):
                if i < len(token_seq) - 1 and token_seq[i] == a and token_seq[i+1] == b:
                    new_seq.append(new_id)
                    i += 2
                else:
                    new_seq.append(token_seq[i])
                    i += 1
            new_tokens.append(new_seq)
        return new_tokens

the _get_stats fn() gets the pairs in a chunck and count its freq.The _merge_pair squishes the freq occured pair by giving a new ID to it.

THE TRAINING

def train(self, text_iterator, pretokenize_pattern=None):
        self.byte_to_char = self._bytes_to_unicode()
        self.char_to_byte = {v: k for k, v in self.byte_to_char.items()}

        for i in range(256):
            char = self.byte_to_char[i]
            self.vocab[char] = i
            self.id_to_token[i] = char

        if pretokenize_pattern is None:
            pretokenize_pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
        pretokenizer = re.compile(pretokenize_pattern)

        all_tokens = []
        for text in text_iterator:
            chunks = pretokenizer.findall(text)
            for chunk in chunks:
                token_ids = list(chunk.encode('utf-8'))
                all_tokens.append(token_ids)

        next_id = 256
        while next_id < self.vocab_size:
            pair_freqs = self._get_stats(all_tokens)
            if not pair_freqs: break
            (a, b), _ = pair_freqs.most_common(1)[0]

            new_token_str = self.id_to_token[a] + self.id_to_token[b]
            self.merges[(a, b)] = next_id
            self.vocab[new_token_str] = next_id
            self.id_to_token[next_id] = new_token_str

            all_tokens = self._merge_pair(a, b, next_id, all_tokens)
            next_id += 1
            if next_id % 1000 == 0:
                print(f"Trained {next_id}/{self.vocab_size} tokens")

The train fn() first takes the mapping of bytes to characters and the reverse of it. Then it processes the text by splitting it into chunks. For each chunk, the chuck IDs are appended to a full list all_tokens. From this list, we find all neighbouring pairs{_get_stats}and get the most frequently occurring one{freqs.most_common(1)[0]} and merge it{_merge_pair}.

The new ID is the ID we assign to every new token or new combination of word we see.

round 1:  "th" → 256    ← super common in English
round 2:  "he" → 257
round 3:  "in" → 258
...
round 500:  "the" → 756    ← "th" + "e", both already merged before
round 2000: " the" → 2256  ← space + "the"
round 8000: " world" → 8256