My first language model, from scratch

1. Give each character a number

Start with a small piece of text: hello world. A tokenizer translates text into a sequence of token IDs. For this first version, one character is one token. Spaces count too.

set(text) collects the unique characters. sorted puts them in a predictable order, and enumerate gives each one an integer starting at zero. The dictionary is our lookup table. Those numbers are labels; a larger ID does not mean a more important character.

text = "hello world"
vocab = sorted(set(text))

print(vocab)

char_to_id = {}

for token_id, character in enumerate(vocab):
    char_to_id[character] = token_id

print(char_to_id)

The vocabulary contains eight characters. The repeated letters in the original text appear only once in the dictionary.

2. Make the round trip

Encoding walks through the text and looks up each character’s ID. Decoding does the reverse: it looks up the character at each position in vocab, then joins those characters into a string.

Add the two functions below the vocabulary. The assertion at the end checks that encoding and decoding together recover the exact original text.

def encode(text):
    token_ids = []

    for character in text:
        token_ids.append(char_to_id[character])

    return token_ids

print(encode("hello"))

def decode(token_ids):
    characters = []

    for token_id in token_ids:
        characters.append(vocab[token_id])

    return "".join(characters)

print(decode([3, 2, 4, 4, 5]))

original = "hello world"
reconstructed = decode(encode(original))

assert reconstructed == original
print("Round trip passed!")

This tokenizer only knows the characters in its vocabulary. Encoding a new character would raise a KeyError. The demo at the end rebuilds the vocabulary from whatever training text you enter.

3. Ask what comes next

Now we need training examples. We take two copies of the token sequence and shift one by a position. Each input token is paired with the token that came immediately after it.

tokens[:-1] excludes the last token, which has no next token in our sample. tokens[1:] excludes the first token. zip walks through those two lists together.

tokens = encode("hello world")

inputs = tokens[:-1]
targets = tokens[1:]

for current_id, next_id in zip(inputs, targets):
    current_character = vocab[current_id]
    next_character = vocab[next_id]

    print(repr(current_character), "->", repr(next_character))

These are the first five pairs. There are ten in total: one fewer than the eleven characters in hello world.

4. Learn by counting

For each current token, keep a dictionary of next tokens and count how often each one appears. This is the learning step in a bigram model. “Bigram” simply means a pair of tokens.

The variable following refers to the inner dictionary. Updating it also updates counts; it is not a separate copy.

counts = {}

for current_id, next_id in zip(inputs, targets):
    if current_id not in counts:
        counts[current_id] = {}

    following = counts[current_id]

    if next_id not in following:
        following[next_id] = 0

    following[next_id] += 1

l_id = char_to_id["l"]

for next_id, count in counts[l_id].items():
    print(repr(vocab[next_id]), count)

The three occurrences of l are followed by l, o, and d. Seeing that list makes the uncertainty concrete: there is more than one possible next character.

5. Turn counts into probabilities

Divide each count by the total count for its starting token. After l, all three continuations have probability 1/3. After h, the only continuation we observed is e, so its probability is 1.

These are probabilities learned from this particular sample. They are not rules of English. For each starting token with observed continuations, the probabilities add up to one, apart from tiny floating-point rounding differences.

probabilities = {}

for current_id, following in counts.items():
    total = sum(following.values())
    probabilities[current_id] = {}

    for next_id, count in following.items():
        probabilities[current_id][next_id] = count / total

for next_id, probability in probabilities[l_id].items():
    print(repr(vocab[next_id]), probability)

No matrix library or machine-learning framework is doing this for us. It is a dictionary, a total, and division.

6. Let it generate

Add import random at the top of the file. It is part of Python’s standard library. We use random.choices to sample a next token according to the probabilities, append it to the output, and repeat.

The function returns a list even when we request one choice, so [0] extracts that choice. Generation stops when it reaches the limit or a character with no observed continuation.

def generate(start_character, max_new_tokens):
    current_id = char_to_id[start_character]
    genereated_ids = [current_id]

    for _ in range(max_new_tokens):
        if current_id not in probabilities:
            break

        following = probabilities[current_id]

        next_id = random.choices(
            population=list(following.keys()),
            weights=list(following.values()),
            k=1,
        )[0]

        genereated_ids.append(next_id)
        current_id = next_id

    return decode(genereated_ids)
print(generate("h", 30))

Try running it several times. Both l → l and l → o are possible, so the model can repeat letters or take a shortcut through the training text. In this sample, d has no continuation, so reaching d stops the run.

The complete working file

This is my original tokenizer.py, including the print statements I used to check each step. Save it and run python tokenizer.py in a terminal. No package installation is needed.

import random

text = "hello world"
vocab = sorted(set(text))

print(vocab)

char_to_id = {}

for token_id, character in enumerate(vocab):
    char_to_id[character] = token_id

print(char_to_id)

def encode(text):
    token_ids = []

    for character in text:
        token_ids.append(char_to_id[character])

    return token_ids

print(encode("hello"))

def decode(token_ids):
    characters = []

    for token_id in token_ids:
        characters.append(vocab[token_id])

    return "".join(characters)

print(decode([3, 2, 4, 4, 5]))

original = "hello world"
reconstructed = decode(encode(original))

assert reconstructed == original
print("Round trip passed!")

tokens = encode("hello world")

inputs = tokens[:-1]
targets = tokens[1:]

for current_id, next_id in zip(inputs, targets):
    current_character = vocab[current_id]
    next_character = vocab[next_id]

    print(repr(current_character), "->", repr(next_character))

counts = {}

for current_id, next_id in zip(inputs, targets):
    if current_id not in counts:
        counts[current_id] = {}

    following = counts[current_id]

    if next_id not in following:
        following[next_id] = 0

    following[next_id] += 1

l_id = char_to_id["l"]

for next_id, count in counts[l_id].items():
    print(repr(vocab[next_id]), count)

probabilities = {}

for current_id, following in counts.items():
    total = sum(following.values())
    probabilities[current_id] = {}

    for next_id, count in following.items():
        probabilities[current_id][next_id] = count / total

for next_id, probability in probabilities[l_id].items():
    print(repr(vocab[next_id]), probability)

def generate(start_character, max_new_tokens):
    current_id = char_to_id[start_character]
    genereated_ids = [current_id]

    for _ in range(max_new_tokens):
        if current_id not in probabilities:
            break

        following = probabilities[current_id]

        next_id = random.choices(
            population=list(following.keys()),
            weights=list(following.values()),
            k=1,
        )[0]

        genereated_ids.append(next_id)
        current_id = next_id

    return decode(genereated_ids)
print(generate("h", 30))