Week 9 — Tokenizers: from raw text to training batches

import torch

text = "我喜欢AI，AI也喜欢猫。"
character_tokens = ["我", "喜", "欢", "A", "I", "，", "也", "猫", "。"]
char_to_id = {token: index for index, token in enumerate(character_tokens)}
id_to_char = {index: token for token, index in char_to_id.items()}


def encode_characters(value: str) -> list[int]:
    return [char_to_id[character] for character in value]


def decode_characters(ids: list[int]) -> str:
    return "".join(id_to_char[index] for index in ids)


ids = encode_characters(text)
token_tensor = torch.tensor(ids, dtype=torch.long)

assert ids == [0, 1, 2, 3, 4, 5, 3, 4, 6, 1, 2, 7, 8]
assert tuple(token_tensor.shape) == (13,)
assert decode_characters(ids) == text

TOKENS = (
    "<BOS>",
    "<EOS>",
    "<PAD>",
    "<UNK>",
    "我",
    "喜欢",
    "AI",
    "，",
    "也",
    "猫",
    "。",
)
TOKEN_TO_ID = {token: token_id for token_id, token in enumerate(TOKENS)}
ID_TO_TOKEN = {token_id: token for token, token_id in TOKEN_TO_ID.items()}

BOS_ID = TOKEN_TO_ID["<BOS>"]  # 0
EOS_ID = TOKEN_TO_ID["<EOS>"]  # 1
PAD_ID = TOKEN_TO_ID["<PAD>"]  # 2
UNK_ID = TOKEN_TO_ID["<UNK>"]  # 3

# Deterministic longest-first routing for this small teaching artifact.
# Equal-length routes retain this declared order.
CONTENT_ROUTES = ("喜欢", "AI", "我", "，", "也", "猫", "。")
HIDDEN_ON_DECODE_IDS = {BOS_ID, EOS_ID, PAD_ID}


class ReadableTokenizerV1:
    version = "w09-readable-v1"
    normalization = "identity"
    vocabulary = TOKENS

    def segment_content(self, text: str) -> list[str]:
        pieces: list[str] = []
        cursor = 0
        while cursor < len(text):
            matched = next(
                (
                    piece
                    for piece in CONTENT_ROUTES
                    if text.startswith(piece, cursor)
                ),
                None,
            )
            if matched is None:
                pieces.append("<UNK>")
                cursor += 1
            else:
                pieces.append(matched)
                cursor += len(matched)
        return pieces

    def encode_content(
        self,
        text: str,
        *,
        add_special_tokens: bool = False,
    ) -> list[int]:
        if add_special_tokens:
            raise ValueError("encode_content never adds BOS/EOS")
        return [TOKEN_TO_ID[piece] for piece in self.segment_content(text)]

    def encode_document(
        self,
        text: str,
        *,
        add_special_tokens: bool = True,
    ) -> list[int]:
        if not add_special_tokens:
            raise ValueError("use encode_content when boundaries are unwanted")
        return [
            BOS_ID,
            *self.encode_content(text, add_special_tokens=False),
            EOS_ID,
        ]

    def decode(
        self,
        ids: list[int],
        *,
        skip_special_tokens: bool = True,
    ) -> str:
        pieces: list[str] = []
        for token_id in ids:
            if token_id not in ID_TO_TOKEN:
                raise ValueError(f"token ID out of range: {token_id}")
            if skip_special_tokens and token_id in HIDDEN_ON_DECODE_IDS:
                continue
            pieces.append(ID_TO_TOKEN[token_id])
        return "".join(pieces)


W09_READABLE_V1 = ReadableTokenizerV1()


def encode_content(text: str) -> list[int]:
    return W09_READABLE_V1.encode_content(
        text,
        add_special_tokens=False,
    )


def encode_document(text: str) -> list[int]:
    return W09_READABLE_V1.encode_document(
        text,
        add_special_tokens=True,
    )


RUNNING_TEXT = "我喜欢AI，AI也喜欢猫。"
content_ids = encode_content(RUNNING_TEXT)
document_ids = encode_document(RUNNING_TEXT)

assert content_ids == [4, 5, 6, 7, 6, 8, 5, 9, 10]
assert document_ids == [0, *content_ids, 1]
assert W09_READABLE_V1.decode(
    document_ids,
    skip_special_tokens=True,
) == RUNNING_TEXT

# These calls never mutate vocabulary, routes, or IDs.

w09-readable-v1:
ID 5 -> 喜欢
ID 6 -> AI

incompatible-tokenizer:
ID 5 -> AI
ID 6 -> 喜欢

Both report V=11.
Model input ID 6 is numerically in range, but its learned row and displayed token disagree.

import torch

from w09_readable_v1 import encode_document

documents = ["我喜欢AI，AI也喜欢猫。"]

# encode_document already inserts exactly one BOS and one EOS per document.
stream: list[int] = []
for document in documents:
    stream.extend(encode_document(document))

stream_tensor = torch.tensor(stream, dtype=torch.long)

assert stream == [0, 4, 5, 6, 7, 6, 8, 5, 9, 10, 1]
assert tuple(stream_tensor.shape) == (11,)
assert stream_tensor.dtype == torch.long

x = s[0:4] = [0,4,5,6] = [<BOS>, 我, 喜欢, AI]  shape [T]=[4]
y = s[1:5] = [4,5,6,7] = [我, 喜欢, AI, ，]       shape [T]=[4]

starts = [0,5]

inputs = [
  [0,4,5,6],  # <BOS> 我 喜欢 AI
  [6,8,5,9],  # AI 也 喜欢 猫
]             # [B,T] = [2,4]

targets = [
  [4,5,6,7],  # 我 喜欢 AI ，
  [8,5,9,10], # 也 喜欢 猫 。
]             # [B,T] = [2,4]

import torch


def sample_batch(source: torch.Tensor, batch_size: int, context_length: int):
    if source.ndim != 1:
        raise ValueError("source must have shape [L]")
    if len(source) <= context_length:
        raise ValueError("source needs at least context_length + 1 IDs")

    # torch.randint excludes high, so high=L-T yields starts 0..L-T-1.
    starts = torch.randint(
        low=0,
        high=len(source) - context_length,
        size=(batch_size,),
    )
    inputs = torch.stack(
        [source[i : i + context_length] for i in starts.tolist()]
    )
    targets = torch.stack(
        [source[i + 1 : i + context_length + 1] for i in starts.tolist()]
    )
    return inputs, targets


source = torch.tensor([0, 4, 5, 6, 7, 6, 8, 5, 9, 10, 1])
inputs, targets = sample_batch(source, batch_size=2, context_length=4)
assert inputs.shape == targets.shape == (2, 4)

from w09_readable_v1 import (
    BOS_ID,
    EOS_ID,
    W09_READABLE_V1,
    encode_content,
    encode_document,
)

text = "我喜欢AI，AI也喜欢猫。"
content_ids = encode_content(text)
document_ids = encode_document(text)

assert document_ids == [BOS_ID, *content_ids, EOS_ID]
assert W09_READABLE_V1.decode(
    document_ids,
    skip_special_tokens=True,
) == text

# Wrong: encode_document already owns both boundaries.
# duplicated = [BOS_ID, *encode_document(text), EOS_ID]

import torch
import torch.nn.functional as F

input_ids = torch.tensor([
    [0, 4, 5, 6, 7, 6, 8],
    [0, 4, 5, 6, 10, 1, 2],
], dtype=torch.long)

attention_mask = torch.tensor([
    [1, 1, 1, 1, 1, 1, 1],
    [1, 1, 1, 1, 1, 1, 0],
], dtype=torch.bool)

targets = torch.tensor([
    [4, 5, 6, 7, 6, 8, 5],
    [4, 5, 6, 10, 1, -100, -100],
], dtype=torch.long)

B, T = input_ids.shape
causal = torch.tril(torch.ones(T, T, dtype=torch.bool))

# allowed[b,t,j]: query t may read key j only when j is not
# in the future and row b says that key j is real rather than PAD.
allowed = causal.unsqueeze(0) & attention_mask[:, None, :]
assert tuple(allowed.shape) == (2, 7, 7)
assert allowed[1, 4].tolist() == [
    True, True, True, True, True, False, False
]

valid_targets = targets.ne(-100)
assert int(valid_targets.sum()) == 12  # row 0: 7, row 1: 5

# Zero logits mean a uniform distribution over V=11 candidates.
# Cross Entropy averages only the 12 valid targets.
logits = torch.zeros(B, T, 11)
loss = F.cross_entropy(
    logits.reshape(B * T, 11),
    targets.reshape(B * T),
    ignore_index=-100,
)
assert torch.isclose(loss, torch.log(torch.tensor(11.0)))

Round 1
all six candidates tie at count 2
lexicographically smallest integer pair = (0x41, 0x49)

(41,49) -> new token 256, displayed [41 49]

... E6 AC A2 [41 49] EF BC 8C [41 49] E4 B9 9F ...
length: 31 - 2 = 29 symbols

Round 2
all five repeated pairs tie at count 2
lexicographically smallest integer pair = (0x96, 0x9C)

(96,9C) -> new token 257, displayed [96 9C]

喜 at each occurrence is now: E5 [96 9C]
length: 29 - 2 = 27 symbols

total after two rounds: 31 - 2 - 2 = 27

from collections import Counter


def count_pairs(symbols: list[int]) -> Counter:
    return Counter(zip(symbols, symbols[1:]))


def choose_pair(pair_counts: Counter) -> tuple[int, int]:
    highest_count = max(pair_counts.values())
    tied_pairs = [
        pair
        for pair, count in pair_counts.items()
        if count == highest_count
    ]
    return min(tied_pairs)  # lexicographic integer-pair tie break


def merge_non_overlapping(
    symbols: list[int],
    pair: tuple[int, int],
    new_id: int,
) -> tuple[list[int], int]:
    merged: list[int] = []
    replacements = 0
    cursor = 0

    while cursor < len(symbols):
        if (
            cursor + 1 < len(symbols)
            and (symbols[cursor], symbols[cursor + 1]) == pair
        ):
            merged.append(new_id)
            replacements += 1
            cursor += 2
        else:
            merged.append(symbols[cursor])
            cursor += 1

    return merged, replacements


text = "我喜欢AI，AI也喜欢猫。"
training_symbols = list(text.encode("utf-8"))
history: list[tuple[tuple[int, int], int, int, int]] = []

assert len(training_symbols) == 31

for new_id in (256, 257):
    pair_counts = count_pairs(training_symbols)
    chosen_pair = choose_pair(pair_counts)
    before_length = len(training_symbols)
    training_symbols, replacements = merge_non_overlapping(
        training_symbols,
        chosen_pair,
        new_id,
    )
    assert len(training_symbols) == before_length - replacements
    history.append(
        (chosen_pair, new_id, replacements, len(training_symbols))
    )

assert history == [
    ((0x41, 0x49), 256, 2, 29),
    ((0x96, 0x9C), 257, 2, 27),
]


def encode_with_frozen_merges(
    value: str,
    merges: list[tuple[tuple[int, int], int, int, int]],
) -> list[int]:
    symbols = list(value.encode("utf-8"))
    for pair, new_id, _training_replacements, _training_length in merges:
        symbols, _ = merge_non_overlapping(symbols, pair, new_id)
    return symbols


encoded = encode_with_frozen_merges(text, history)
assert encoded == training_symbols
assert len(encoded) == 27

# A learned token stores a byte sequence, so decoding expands first.
token_bytes = {byte_id: bytes([byte_id]) for byte_id in range(256)}
for pair, new_id, _count, _length in history:
    left, right = pair
    token_bytes[new_id] = token_bytes[left] + token_bytes[right]

restored_bytes = b"".join(token_bytes[token_id] for token_id in encoded)
restored_text = restored_bytes.decode("utf-8", errors="strict")
assert restored_text == text

text = "我喜欢AI，AI也喜欢猫。"
utf8_bytes = text.encode("utf-8")

assert len(text) == 13
assert len(utf8_bytes) == 31
assert utf8_bytes.hex(" ").upper() == (
    "E6 88 91 E5 96 9C E6 AC A2 41 49 EF BC 8C "
    "41 49 E4 B9 9F E5 96 9C E6 AC A2 E7 8C AB E3 80 82"
)

import unicodedata

composed = "é"
decomposed = "e\u0301"

assert composed != decomposed
assert [f"U+{ord(ch):04X}" for ch in composed] == ["U+00E9"]
assert [f"U+{ord(ch):04X}" for ch in decomposed] == [
    "U+0065",
    "U+0301",
]
assert composed.encode("utf-8").hex(" ").upper() == "C3 A9"
assert decomposed.encode("utf-8").hex(" ").upper() == "65 CC 81"

# NFC makes these two canonically equivalent spellings identical.
assert unicodedata.normalize("NFC", decomposed) == composed

# NFKC also performs compatibility folding, which is a stronger choice.
assert unicodedata.normalize("NFKC", "ＡＩ") == "AI"

# Standalone: run from the English course_examples directory.
import torch
from course_data import FIVE_WORD_TOKENIZER, DEMO_DOCUMENTS, make_windows

x, y = make_windows(DEMO_DOCUMENTS, FIVE_WORD_TOKENIZER, block_size=2)
inputs = torch.tensor(x, dtype=torch.long)
targets = torch.tensor(y, dtype=torch.long)
print(inputs.tolist())   # [[0,1],[4,1],[0,3]]
print(targets.tolist())  # [[1,2],[1,0],[3,2]]
# Week 10 connects the model: logits, loss = model(inputs, targets)