Токенизатор BPE с нуля: более короткие последовательности для небольшого окна контекста GPT
Напишите токенизатор с кодированием пар байтов на Python, обучите его на WikiText-2, сохраните результаты токенизации в формате JSON и посмотрите, как субсловесная токенизация увеличивает окно контекста GPT.
Источники: Tutorial 16: Beyond Character Tokens, Building a BPE Tokenizer
Токенизатор на уровне символов — это самый простой способ заставить работать небольшую модель GPT: словарный запас минимален, ничего не является неизвестным, и каждый шаг полностью прозрачен. Недостаток заключается в длине последовательности, поскольку каждая буква занимает место в окне контекста. В данном руководстве этот подход заменяется кодированием пар байт (BPE), реализованным с нуля на чистом Python. Вы научите модель работать с данными WikiText-2, сохраните полученный словарный запас и результаты объединения в формате JSON, сгенерируете их «отпечаток», перекодируете набор данных и увидите, какие именно изменения происходят при переобучении модели с использованием субсловесных токенов.
Код предполагает наличие готовой небольшой модели GPT на PyTorch и очищенных файлов WikiText-2. Если при обучении по-прежнему возникают проблемы, сначала стабилизируйте модель; см. тесты для выявления отдельных сбоев в небольшой модели GPT на PyTorch.
Почему символьные токены уже недостаточны
Токенизатор символов присваивает каждому символу, включая пробелы, отдельный токен. Возьмём эту короткую фразу:
artificial intelligence
Он преобразуется в один токен на каждый символ:
['a', 'r', 't', 'i', 'f', 'i', 'c', 'i', 'a', 'l',
' ', 'i', 'n', 't', 'e', 'l', 'l', 'i', 'g', 'e',
'n', 'c', 'e']
Это 23 токена для двух слов. Токенизатор субслов может представить ту же фразу гораздо меньшим количеством элементов, например:
['art', 'ificial', ' intelligence']
Три токена вместо 23. Более короткие последовательности означают больше текста в окне контекста, меньше позиций для обучения и меньше шагов декодирования во время генерации. Модель также получает осмысленные фрагменты вместо изолированных букв, часто встречающиеся фрагменты становятся повторно используемыми единицами, а редкие слова всё ещё можно восстановить из более мелких частей. Сложность заключается в автоматическом поиске хороших субсловесных единиц, и именно это и делает алгоритм BPE.
Сравнение символов, слов и субслов
Существует три основных способа разбиения текста, каждый из которых предполагает определённый компромисс.
Токенизация по символам
Каждый символ является токеном:
"learning"
→ ["l", "e", "a", "r", "n", "i", "n", "g"]
Он предлагает небольшой словарный запас, отсутствие неизвестных токенов и простую реализацию. Взамен вы получаете длинные последовательности, медленную генерацию и модель, которая должна составлять слова из отдельных букв.
Токенизация слов
Каждое слово является токеном:
"machine learning works"
→ ["machine", "learning", "works"]
Последовательности короткие и легко читаемые, но словарный запас становится огромным; для неизвестных слов требуется отдельный токен, а связанные формы, такие как "learn", "learning" и "learned", могут превратиться в три независимых ID без какой-либо общей особенности. Имена, опечатки и редкие слова обрабатываются плохо.
Субсловесная токенизация
Слова делятся на часто встречающиеся фрагменты:
"unbelievable"
→ ["un", "believ", "able"]
Это позволяет получать более короткие последовательности, чем при использовании символов, и меньший словарь, чем при использовании слов. Часто встречающиеся слова могут оставаться отдельными токенами, редкие слова разлагаются на известные составляющие, а фрагменты, такие как префиксы и суффиксы, повторно используются во многих словах. Единственная реальная проблема заключается в том, что обучение и реализация токенизатора требуют больше усилий.
Можно рассматривать эти три подхода как точки на спектре:
Characters
← flexibility
Subwords
← balance
Words
← compactness
Как BPE учится выполнять слияния
BPE начинает с самых мелких единиц — здесь это отдельные символы — и постоянно сливает наиболее частую пару соседних символов в новый токен. Рассмотрим небольшой корпус из трех слов:
low
lower
lowest
Изначально каждое слово представляет собой просто его символы:
l o w
l o w e r
l o w e s t
Далее укажите пары соседних символов, которые встречаются:
(l, o)
(o, w)
(w, e)
(e, r)
(e, s)
(s, t)
Одна пара выделяется:
(l, o)
Он встречается во всех трех словах, поэтому он объединяется в один символ:
(l, o) → lo
После замены его во всех местах корпус выглядит так:
lo w
lo w e r
lo w e s t
При повторном подсчёте наиболее частой парой теперь является:
(lo, w)
Он также появляется три раза, поэтому его тоже объединяют:
(lo, w) → low
И корпус становится:
low
low e r
low e s t
После достаточного количества раундов наиболее частые последовательности становятся отдельными токенами.
Что даёт цикл обучения
Процедура, которую вы собираетесь реализовать, включает следующие шаги:
Training text
→ initialize character tokens
→ count adjacent token pairs
→ select most frequent pair
→ merge the selected pair
→ add merged token to vocabulary
→ repeat until target vocabulary size
Он генерирует два результата, и оба необходимы для дальнейшего использования токенизатора.
Словарь
Словарь соотносит каждую строку токена с целочисленным идентификатором, например:
"a" → 0
"b" → 1
"th" → 98
"ing" → 127
" the" → 204
Правила слияния
Правила слияния — это упорядоченный список пар, которые были объединены во время обучения:
("t", "h") → "th"
("i", "n") → "in"
("in", "g") → "ing"
(" ", "the") → " the"
Порядок имеет решающее значение. Более позднее правило, такое как ("in", "g"), может сработать только тогда, когда более раннее правило, создавшее "in", уже было применено. Применение тех же правил в другом порядке приводит к различной сегментации.
Первая реализация на простом корпусе данных
Прежде чем приступать к работе с WikiText-2, опробуйте алгоритм на строке, которую можно визуально проверить:
text = "low lower lowest"
Разделение на символы даёт список начальных токенов:
tokens = list(text)
print(tokens)
При печати отображаются все символы, включая пробелы:
['l', 'o', 'w', ' ', 'l', 'o', 'w', 'e', 'r',
' ', 'l', 'o', 'w', 'e', 's', 't']
Чтобы найти кандидатов для слияния, подсчитайте каждую пару соседних элементов. Слияние списка с самим собой с сдвигом на одну позицию позволяет получить каждую соседнюю пару, а класс Counter подсчитывает их:
from collections import Counter
def count_pairs(tokens):
return Counter(
zip(tokens[:-1], tokens[1:])
)
Проверьте наиболее распространённые пары:
pair_counts = count_pairs(tokens)
print(pair_counts.most_common(10))
Результат будет выглядеть примерно так:
[(('l', 'o'), 3),
(('o', 'w'), 3),
(('w', 'e'), 2),
...]
Обратите внимание, что (' ', 'l') также фигурирует в статистике. Поскольку пробелы являются обычными символами в этой схеме, алгоритм BPE может обучаться токенам, содержащим пробел в начале, и именно так позже возникают такие токены, как " the".
Слияние одной пары
Основная операция проходит по списку токенов слева направо. Каждый раз, когда текущий токен и следующий соответствуют выбранной паре, генерируется слиянный токен, и процесс проскальзывает на две позиции вперед; в противном случае копируется текущий токен, и процесс продолжается на одну позицию. Накладки разрешаются слева направо: в случае a a a с парой (a, a) сливаются только первые два символа.
def merge_pair(
tokens,
pair,
merged_token,
):
output = []
index = 0
while index < len(tokens):
if (
index < len(tokens) - 1
and tokens[index] == pair[0]
and tokens[index + 1] == pair[1]
):
output.append(merged_token)
index += 2
else:
output.append(tokens[index])
index += 1
return output
Используйте его для объединения l и o:
tokens = merge_pair(
tokens=tokens,
pair=("l", "o"),
merged_token="lo",
)
print(tokens)
Результат короче, чем входной текст:
['lo', 'w', ' ', 'lo', 'w', 'e', 'r',
' ', 'lo', 'w', 'e', 's', 't']
При каждом слиянии удаляется по одному токену за каждое вхождение пары, и именно отсюда и происходит сжатие.
Повторение до заполнения словаря
Минимальный тренер просто циклически выполняет следующие действия: подсчитывает пары, выбирает наиболее часто встречающуюся, сливает её во всех местах, записывает новый токен и правило, а затем останавливается, когда размер словаря достигает заданного значения. Он также останавливается раньше, если пар не осталось или если наиболее частая пара встречается менее двух раз, поскольку в таком случае слияние больше не создаёт ничего, что можно было бы использовать повторно.
def train_bpe_on_tokens(
tokens,
target_vocab_size,
):
vocabulary = set(tokens)
merges = []
while len(vocabulary) < target_vocab_size:
pair_counts = count_pairs(tokens)
if not pair_counts:
break
best_pair, frequency = (
pair_counts.most_common(1)[0]
)
if frequency < 2:
break
merged_token = "".join(best_pair)
tokens = merge_pair(
tokens=tokens,
pair=best_pair,
merged_token=merged_token,
)
vocabulary.add(merged_token)
merges.append(best_pair)
print(
f"Merged {best_pair} "
f"into {merged_token!r}, "
f"frequency={frequency}"
)
return tokens, vocabulary, merges
Достаточно запустить его на простой модели с небольшой целью, чтобы наблюдать, как происходят слияния по одному:
final_tokens, vocabulary, merges = (
train_bpe_on_tokens(
tokens=list(text),
target_vocab_size=20,
)
)
Эта версия хороша для формирования интуиции, но при каждой итерации происходит повторный анализ всего списка в Python. На полном корпусе WikiText это делает процесс очень медленным, поэтому в практической версии ниже обучение проводится на настраиваемом фрагменте WikiText-2.
Проектирование многоразового токенизатора
Класс токенизатора должен обучаться на необработанном тексте, начиная с отдельных символов, выполнять кодирование и декодирование, отклонять символы, отсутствующие в его словаре, сохранять и загружать данные в формате JSON, а также вычислять «отпечаток» данных. Вокруг него скрипт перекодирует каждый фрагмент WikiText-2 и показывает, насколько короче становятся последовательности BPE.
Следует четко оговорить одно из решений по дизайну. Данный токенизатор работает с символами Unicode, а не с сырыми байтами. Токенизаторы, используемые в крупных моделях в производстве, обычно работают на уровне байтов, что означает, что их базовый словарь состоит из 256 возможных значений байтов, и никакой входной данных не может остаться нерепрезентированным. Реализованная здесь версия, основанная на символах, следует тому же алгоритму и более удобна для чтения, но она не является прямым аналогом этих токенизаторов.
Костяк класса
Конструктор принимает базовые символы и список для их объединения в определенном порядке. Он формирует полный список токенов, начиная с базовых символов и добавляя результаты каждого объединения, пропуская дубликаты. На основе этого списка создаются функции преобразования stoi (строка в ID) и itos (ID в строку), а также фиксируется размер словаря. Поскольку ID формируются на основе позиций в списке, одни и те же базовые токены и операции объединения всегда дают одинаковые ID.
class BPETokenizer:
def __init__(
self,
base_tokens,
merges,
):
self.base_tokens = list(
base_tokens
)
self.merges = [
tuple(pair)
for pair in merges
]
self.tokens = list(
self.base_tokens
)
for left, right in self.merges:
merged = left + right
if merged not in self.tokens:
self.tokens.append(merged)
self.stoi = {
token: index
for index, token in enumerate(
self.tokens
)
}
self.itos = {
index: token
for token, index
in self.stoi.items()
}
self.vocab_size = len(
self.tokens
)
Словарный запас, полученный в результате, состоит следовательно из двух частей:
Base character tokens
+
learned merged tokens
Обучение на реальном тексте
Метод обучения превращается в классовый метод, возвращающий готовый к использованию токенизатор. Он сортирует уникальные символы текста для формирования базового словаря, отклоняет целевой размер, меньший, чем этот базовый набор, а затем выполняет тот же цикл подсчёта, выбора и слияния, что и ранее. Логирование ограничивается: выводятся первые 20 операций слияния, а затем каждая стоящая после них, вместе с частотой пар и текущей длиной последовательности, чтобы можно было отслеживать прогресс сжатия, не теряясь в большом объёме данных.
@classmethod
def train(
cls,
text,
target_vocab_size,
minimum_pair_frequency=2,
verbose=True,
):
base_tokens = sorted(
set(text)
)
if target_vocab_size < len(base_tokens):
raise ValueError(
"target_vocab_size cannot be "
"smaller than the number of "
"unique characters"
)
working_tokens = list(text)
vocabulary = set(base_tokens)
merges = []
while len(vocabulary) < target_vocab_size:
pair_counts = Counter(
zip(
working_tokens[:-1],
working_tokens[1:],
)
)
if not pair_counts:
break
best_pair, frequency = (
pair_counts.most_common(1)[0]
)
if frequency < minimum_pair_frequency:
break
merged_token = "".join(
best_pair
)
working_tokens = merge_pair(
tokens=working_tokens,
pair=best_pair,
merged_token=merged_token,
)
merges.append(best_pair)
vocabulary.add(merged_token)
if (
verbose
and (
len(merges) <= 20
or len(merges) % 100 == 0
)
):
print(
f"Merge {len(merges):4d}: "
f"{best_pair!r} "
f"→ {merged_token!r}, "
f"frequency={frequency:,}, "
f"tokens={len(working_tokens):,}"
)
return cls(
base_tokens=base_tokens,
merges=merges,
)
Почему существует минимальная частота пар
Если наиболее частая оставшаяся пара встречается только один раз, её объединение приводит к созданию токена, который никогда не будет использован снова, что просто тратит слот в словаре. Поэтому цикл останавливается при таких условиях:
frequency < minimum_pair_frequency
Для подобных экспериментов достаточен порог в два:
minimum_pair_frequency = 2
При работе с более крупным корпусом текста может потребоваться более высокий порог, чтобы словарь использовался только для тех шаблонов, которые встречаются достаточно часто, чтобы модель могла выучить для них качественные эмбеддинги.
Кодирование и декодирование с помощью токенизатора
Процесс кодирования включает три шага:
- Разбить текст на базовые символы.
- Выполнить все правила объединения в том порядке, в котором они были выучены.
- Найти ID каждого полученного токена.
Первая половина этой работы реализована в функции encode_tokens. Сначала собираются все символы, отсутствующие в словаре, и при их наличии выдается четкая ошибка, что гораздо проще для отладки, чем возникновение ошибки KeyError в процессе обработки набора данных. Затем выполняются операции объединения.
def encode_tokens(self, text):
unknown_characters = sorted({
char
for char in text
if char not in self.stoi
})
if unknown_characters:
raise ValueError(
"Text contains characters missing "
"from the base vocabulary: "
f"{unknown_characters}"
)
tokens = list(text)
for pair in self.merges:
merged_token = (
pair[0] + pair[1]
)
tokens = merge_pair(
tokens=tokens,
pair=pair,
merged_token=merged_token,
)
return tokens
Преобразование полученных строк токенов в ID — это простой поиск по таблице:
def encode(self, text):
tokens = self.encode_tokens(text)
return [
self.stoi[token]
for token in tokens
]
Декодирование ещё проще. Каждый токен, объединённый или нет, хранит точный текст, который он охватывает, поэтому декодирование сводится к простой конкатенации:
def decode(self, token_ids):
return "".join(
self.itos[int(token_id)]
for token_id in token_ids
)
Это придает BPE свойство, которое стоит явно проверить — без потерь при обратном преобразовании:
decode(encode(text)) = text
Быстрая проверка с примерным предложением:
sample = "The history of science."
encoded = tokenizer.encode(sample)
decoded = tokenizer.decode(encoded)
assert decoded == sample
Сохранение и перезагрузка токенизатора
Вам не нужно хранить сам словарь, поскольку его можно детерминированно воссоздать на основе базовых токенов и списка объединения. Метод save записывает эти два элемента в формат JSON вместе с меткой типа и размером словаря для удобства. Установка параметра ensure_ascii=False позволяет сохранять некодируемые символы в файле в читаемом виде, а также создает родительский каталог, если он ещё не существует.
def save(self, path):
data = {
"type": "character_bpe",
"base_tokens": self.base_tokens,
"merges": [
[left, right]
for left, right in self.merges
],
"vocab_size": self.vocab_size,
}
path = Path(path)
path.parent.mkdir(
parents=True,
exist_ok=True,
)
with open(
path,
"w",
encoding="utf-8",
) as file:
json.dump(
data,
file,
ensure_ascii=False,
indent=2,
)
Результатом является один файл:
tokenizer/
└── bpe_tokenizer.json
Загрузка обратно выполняет процесс в обратном порядке, считывая JSON и передавая обе списки конструктору:
@classmethod
def from_file(cls, path):
with open(
path,
"r",
encoding="utf-8",
) as file:
data = json.load(file)
return cls(
base_tokens=data[
"base_tokens"
],
merges=data["merges"],
)
Затем убедитесь, что перезагруженный токенизатор по-прежнему обрабатывает пример текста:
loaded_tokenizer = (
BPETokenizer.from_file(
"tokenizer/bpe_tokenizer.json"
)
)
assert (
loaded_tokenizer.decode(
loaded_tokenizer.encode(sample)
)
== sample
)
Обучение на фрагменте WikiText-2
Поскольку простой тренер пересканирует всю последовательность данных при каждом слиянии, обучение всего объема WikiText-2 с большим словарным запасом может занять много времени. Начните с настраиваемого префикса текста для обучения:
training_characters = 500_000
training_text = train_text[
:training_characters
]
Разумный размер начального словаря составляет:
target_vocab_size = 500
Другими словами:
Initial character vocabulary
+
learned merged tokens
=
approximately 500 tokens
Когда всё работает, попробуйте более крупные цели:
target_vocab_size = 1_000
target_vocab_size = 2_000
Ожидайте, что время обучения будет увеличиваться как с ростом размера выборки, так и с увеличением количества слияний.
Использование базовых символов из каждого разделения
Если базовый словарь формируется исключительно на основе обучающей выборки, в тексте для проверки или тестирования могут встречаться символы, которых там никогда не было, и кодирование таких разделений завершится ошибкой. Решение заключается в сборе базового набора символов из всех трех разделений, при этом статистику слияний продолжая считать только на основе обучающих данных:
all_text = (
train_text
+ val_text
+ test_text
)
base_tokens = sorted(
set(all_text)
)
Это не приводит к просачиванию каких-либо шаблонов языка валидации в результаты слияния. Оно лишь гарантирует, что у каждого символа в наборе данных есть идентификатор. Для поддержки этого трейнер получает необязательный аргумент base_tokens, который используется вместо символов учебного текста, если он указан:
@classmethod
def train(
cls,
text,
target_vocab_size,
minimum_pair_frequency=2,
base_tokens=None,
verbose=True,
):
if base_tokens is None:
base_tokens = sorted(
set(text)
)
else:
base_tokens = sorted(
set(base_tokens)
)
Определение степени сжатия
Чтобы понять, какую пользу приносит BPE, сравните количество токенов в одном и том же отрывке. Возьмем пример:
sample_text = train_text[:10_000]
Количество токенов в его строке равно просто её длине:
character_count = len(sample_text)
Количество токенов BPE равно длине списка закодированных идентификаторов:
bpe_ids = tokenizer.encode(
sample_text
)
bpe_count = len(bpe_ids)
Соотношение между ними — это коэффициент сжатия:
compression_factor = (
character_count / bpe_count
)
Результат может выглядеть так:
Character tokens: 10,000
BPE tokens: 4,250
Compression: 2.35×
При коэффициенте сжатия 2,35 блок из 128 токенов BPE покрывает примерно столько же текста, сколько 301 токен, состоящий из символов. Однако сжатие не является однородным: проза, наполненная распространёнными словами, сжимается хорошо, в то время как числа, редкие имена и необычная пунктуация сжимаются плохо.
Что выучил токенизатор
Выученные токены — это все элементы словаря, находящиеся после базовых символов. Если отсортировать их по длине и вывести самые длинные из них вместе с их идентификаторами, это позволит быстро понять, что токенизатор считает частым:
learned_tokens = (
tokenizer.tokens[
len(tokenizer.base_tokens):
]
)
longest_tokens = sorted(
learned_tokens,
key=len,
reverse=True,
)
for token in longest_tokens[:50]:
print(
repr(token),
tokenizer.stoi[token],
)
Обычно вы увидите токены примерно в таком виде:
" the"
"ing"
"tion"
" of"
"ed"
" and"
Точный список зависит от обучающей выборки, предобработки, размера целевого словаря, минимальной частоты и правил разрешения равенства. Что касается правил разрешения равенства: функция Counter.most_common возвращает пары с одинаковой частотой в том порядке, в котором они были впервые подсчитаны, поэтому результаты для данного входного данных являются детерминированными, но изменение выборки может изменить порядок, в котором победят среди нескольких пар с равной частотой.
Перекодирование набора данных
С использованием обученного токенизатора необходимо закодировать все три разделения данных:
train_bpe_ids = tokenizer.encode(
train_text
)
val_bpe_ids = tokenizer.encode(
val_text
)
test_bpe_ids = tokenizer.encode(
test_text
)
Преобразуйте списки ID в тензоры типа long — это тип, который ожидают слои встраивания для индексов:
train_tensor = torch.tensor(
train_bpe_ids,
dtype=torch.long,
)
val_tensor = torch.tensor(
val_bpe_ids,
dtype=torch.long,
)
test_tensor = torch.tensor(
test_bpe_ids,
dtype=torch.long,
)
Сохраните их в отдельный каталог:
data/
└── encoded_bpe/
├── train_ids.pt
├── val_ids.pt
└── test_ids.pt
Не трогайте файлы с кодированными символами; сохранение обоих вариантов облегчает последующее сравнение и предотвращает передачу ID одной модели другой.
Проверка кодированных разделов
Во-первых, каждый ID должен находиться в диапазоне от нуля до размера словаря включительно, но без учета самого размера словаря:
for split_name, token_ids in [
("train", train_tensor),
("validation", val_tensor),
("test", test_tensor),
]:
assert token_ids.min().item() >= 0
assert (
token_ids.max().item()
< tokenizer.vocab_size
)
print(
split_name,
token_ids.shape,
token_ids.min().item(),
token_ids.max().item(),
)
Во-вторых, пример из каждой группы должен возвращаться к исходному тексту после обработки:
for text in [
train_text[:1_000],
val_text[:1_000],
test_text[:1_000],
]:
encoded = tokenizer.encode(text)
decoded = tokenizer.decode(encoded)
assert decoded == text
Интеграция BPE в модель
Архитектура модели не требует структурных изменений. Изменяется лишь размер словаря в настройках:
config = MiniGPTConfig(
vocab_size=bpe_tokenizer.vocab_size,
block_size=128,
embedding_dim=128,
num_heads=4,
num_layers=4,
expansion_factor=4,
dropout=0.1,
)
Направьте загрузку данных на новые файлы:
train_ids = torch.load(
"data/encoded_bpe/train_ids.pt"
)
Начиная с этого момента модель предсказывает токены BPE вместо символов.
Старые чекпоинты не будут загружаться
У модели, работающей с символами, словарь обычно находится в следующем диапазоне:
100 to 300 character tokens
Модель BPE обычно имеет примерно следующий вид:
500 to 2,000 subword tokens
И вживление токенов, и выходной блок зависят от размера словаря, поэтому их структуры меняются. У контрольной точки для символов есть:
token embedding: [character_vocab_size, C]
В то время как модель BPE ожидает:
token embedding: [bpe_vocab_size, C]
Здесь C — это размерность вкладки. Формы не совпадают, поэтому старые веса нельзя загрузить напрямую. Даже если бы это было возможно, ID имеют разное значение в разных токенизаторах. Вам необходимо обучить новую модель, и полезно хранить чекпоинты отдельно:
checkpoints_char/
checkpoints_bpe/
Или, что эквивалентно:
checkpoints/char/
checkpoints/bpe/
Отпечаток идентификатора токенизатора
Различия в форме токенов заметны; однако сопоставление чекпоинта с другим токенизатором BPE той же размерности остается незаметным, поскольку каждый идентификатор просто соответствует неверному тексту. Функция отпечатка генерирует хэш в компактном JSON-формате базовых токенов и объединяет его с SHA-256, так что любое изменение приводит к появлению нового значения.
import hashlib
def tokenizer_fingerprint(tokenizer):
payload = json.dumps(
{
"base_tokens": (
tokenizer.base_tokens
),
"merges": [
list(pair)
for pair in tokenizer.merges
],
},
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
Храните отпечаток пальца внутри каждой точки контроля:
checkpoint[
"tokenizer_fingerprint"
] = tokenizer_fingerprint(
tokenizer
)
Во время загрузки вычисляется «отпечаток» действующего токенизатора, и если он не совпадает с тем, который сохранён в чекпоинте, выполнение прерывается.
Ограничение производительности, о котором стоит знать
Кодировщик переисполняет каждое правило объединения для всей последовательности токенов:
for pair in self.merges:
tokens = merge_pair(...)
Это означает один полный проход по тексту при каждом слиянии, поэтому большие тексты с тысячами слияний кодируются медленно. Реализация преднамеренно отдает приоритет ясности, корректности и возможности проверки над скоростью. Производственные токенизаторы заранее делят текст на слова, хранят в кэше повторяющиеся разбиения и применяют слияния в определенном порядке; для обучения же предпочтительна прозрачная версия.
Полная программа
Соберите всё в один скрипт:
build_bpe_tokenizer.py
Полная программа ниже объединяет в себе инструменты для работы с парами, класс токенизатора, функции проверки корректности данных и точку входа в командную строку. Два аспекта отличаются от приведённых примеров: при сохранении теперь сохраняется «отпечаток» данных, и запись происходит через временный файл, который затем заменяется целевым файлом, так что прерванный запуск никогда не оставляет незавершённого файла; кроме того, в процессе обучения проверяется, что базовый словарь покрывает текст для обучения. Основная функция обучает модель на настроенном префиксе, сохраняет и перезагружает токенизатор, проверяет «отпечатки» данных, выполняет двусторонние операции и сохраняет каждый раздел данных, проверяет диапазоны идентификаторов и выводит статистику.
Похоже, что одна строка в цикле «самых длинных изученных токенов» была повреждена: f"{loaded_tokenizer.stoi5d} " вызовет ошибку AttributeError, поскольку функция stoi5d не существует. Предполагаемое выражение, скорее всего, — f"{loaded_tokenizer.stoi[token]:5d} ", которое выводит ID каждого токена с заполнением до пяти символов. Исправьте эту строку перед запуском скрипта.
import argparse
import hashlib
import json
from collections import Counter
from pathlib import Path
import torch
# ==================================================
# Pair Operations
# ==================================================
def count_pairs(tokens):
return Counter(
zip(tokens[:-1], tokens[1:])
)
def merge_pair(
tokens,
pair,
merged_token,
):
output = []
index = 0
while index < len(tokens):
if (
index < len(tokens) - 1
and tokens[index] == pair[0]
and tokens[index + 1] == pair[1]
):
output.append(merged_token)
index += 2
else:
output.append(tokens[index])
index += 1
return output
# ==================================================
# BPE Tokenizer
# ==================================================
class BPETokenizer:
def __init__(
self,
base_tokens,
merges,
):
self.base_tokens = list(
base_tokens
)
self.merges = [
tuple(pair)
for pair in merges
]
self.tokens = list(
self.base_tokens
)
for left, right in self.merges:
merged_token = left + right
if merged_token not in self.tokens:
self.tokens.append(
merged_token
)
self.stoi = {
token: index
for index, token
in enumerate(self.tokens)
}
self.itos = {
index: token
for token, index
in self.stoi.items()
}
self.vocab_size = len(
self.tokens
)
@classmethod
def train(
cls,
text,
target_vocab_size,
minimum_pair_frequency=2,
base_tokens=None,
verbose=True,
):
if not text:
raise ValueError(
"Training text cannot be empty"
)
if base_tokens is None:
base_tokens = sorted(
set(text)
)
else:
base_tokens = sorted(
set(base_tokens)
)
missing_characters = (
set(text) - set(base_tokens)
)
if missing_characters:
raise ValueError(
"Base vocabulary is missing "
"training characters"
)
if (
target_vocab_size
< len(base_tokens)
):
raise ValueError(
"target_vocab_size cannot be "
"smaller than the number of "
"base tokens"
)
working_tokens = list(text)
vocabulary = set(base_tokens)
merges = []
print(
"Initial characters:",
len(working_tokens),
)
print(
"Base vocabulary size:",
len(base_tokens),
)
while (
len(vocabulary)
< target_vocab_size
):
pair_counts = count_pairs(
working_tokens
)
if not pair_counts:
print(
"No adjacent pairs remain."
)
break
best_pair, frequency = (
pair_counts.most_common(1)[0]
)
if (
frequency
< minimum_pair_frequency
):
print(
"Stopping because the most "
"frequent pair occurs only "
f"{frequency} time(s)."
)
break
merged_token = "".join(
best_pair
)
working_tokens = merge_pair(
tokens=working_tokens,
pair=best_pair,
merged_token=merged_token,
)
merges.append(best_pair)
vocabulary.add(merged_token)
merge_number = len(merges)
if (
verbose
and (
merge_number <= 20
or merge_number % 100 == 0
)
):
print(
f"Merge {merge_number:4d}: "
f"{best_pair!r} "
f"→ {merged_token!r}, "
f"frequency={frequency:,}, "
f"sequence length="
f"{len(working_tokens):,}, "
f"vocabulary="
f"{len(vocabulary):,}"
)
tokenizer = cls(
base_tokens=base_tokens,
merges=merges,
)
print(
"Final vocabulary size:",
tokenizer.vocab_size,
)
print(
"Final training sequence length:",
len(working_tokens),
)
return tokenizer
@classmethod
def from_file(cls, path):
with open(
path,
"r",
encoding="utf-8",
) as file:
data = json.load(file)
return cls(
base_tokens=data[
"base_tokens"
],
merges=data["merges"],
)
def encode_tokens(self, text):
unknown_characters = sorted({
char
for char in text
if char not in self.stoi
})
if unknown_characters:
raise ValueError(
"Text contains characters "
"missing from the base vocabulary: "
f"{unknown_characters}"
)
tokens = list(text)
for left, right in self.merges:
tokens = merge_pair(
tokens=tokens,
pair=(left, right),
merged_token=left + right,
)
return tokens
def encode(self, text):
return [
self.stoi[token]
for token in self.encode_tokens(
text
)
]
def decode(self, token_ids):
return "".join(
self.itos[int(token_id)]
for token_id in token_ids
)
def save(self, path):
path = Path(path)
path.parent.mkdir(
parents=True,
exist_ok=True,
)
data = {
"type": "character_bpe",
"base_tokens": (
self.base_tokens
),
"merges": [
[left, right]
for left, right
in self.merges
],
"vocab_size": (
self.vocab_size
),
"fingerprint": (
self.fingerprint()
),
}
temporary_path = path.with_suffix(
".tmp"
)
with open(
temporary_path,
"w",
encoding="utf-8",
) as file:
json.dump(
data,
file,
ensure_ascii=False,
indent=2,
)
temporary_path.replace(path)
def fingerprint(self):
payload = json.dumps(
{
"base_tokens": (
self.base_tokens
),
"merges": [
list(pair)
for pair in self.merges
],
},
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(
payload
).hexdigest()
# ==================================================
# Validation
# ==================================================
def validate_round_trip(
tokenizer,
text,
split_name,
):
encoded = tokenizer.encode(text)
decoded = tokenizer.decode(encoded)
if decoded != text:
raise AssertionError(
f"Round-trip failed for "
f"{split_name}"
)
print(
f"{split_name} round-trip passed: "
f"{len(text):,} characters "
f"→ {len(encoded):,} tokens"
)
return encoded
def save_token_tensor(
token_ids,
path,
):
path = Path(path)
path.parent.mkdir(
parents=True,
exist_ok=True,
)
tensor = torch.tensor(
token_ids,
dtype=torch.long,
)
temporary_path = path.with_suffix(
".tmp"
)
torch.save(
tensor,
temporary_path,
)
temporary_path.replace(path)
return tensor
# ==================================================
# Arguments
# ==================================================
def parse_args():
parser = argparse.ArgumentParser(
description=(
"Train a simple educational BPE "
"tokenizer on WikiText-2"
)
)
parser.add_argument(
"--processed-dir",
type=str,
default="data/processed",
)
parser.add_argument(
"--tokenizer-output",
type=str,
default=(
"tokenizer/bpe_tokenizer.json"
),
)
parser.add_argument(
"--encoded-output-dir",
type=str,
default="data/encoded_bpe",
)
parser.add_argument(
"--target-vocab-size",
type=int,
default=500,
)
parser.add_argument(
"--training-characters",
type=int,
default=500_000,
)
parser.add_argument(
"--minimum-pair-frequency",
type=int,
default=2,
)
return parser.parse_args()
# ==================================================
# Main Program
# ==================================================
def main():
args = parse_args()
processed_dir = Path(
args.processed_dir
)
train_path = (
processed_dir
/ "wikitext2_train_clean.txt"
)
val_path = (
processed_dir
/ "wikitext2_val_clean.txt"
)
test_path = (
processed_dir
/ "wikitext2_test_clean.txt"
)
for path in [
train_path,
val_path,
test_path,
]:
if not path.exists():
raise FileNotFoundError(
f"Required file not found: "
f"{path}"
)
train_text = train_path.read_text(
encoding="utf-8"
)
val_text = val_path.read_text(
encoding="utf-8"
)
test_text = test_path.read_text(
encoding="utf-8"
)
print(
"Training characters:",
len(train_text),
)
print(
"Validation characters:",
len(val_text),
)
print(
"Test characters:",
len(test_text),
)
all_characters = sorted(
set(
train_text
+ val_text
+ test_text
)
)
training_character_count = min(
args.training_characters,
len(train_text),
)
bpe_training_text = train_text[
:training_character_count
]
print(
"\nBPE training characters:",
len(bpe_training_text),
)
tokenizer = BPETokenizer.train(
text=bpe_training_text,
target_vocab_size=(
args.target_vocab_size
),
minimum_pair_frequency=(
args.minimum_pair_frequency
),
base_tokens=all_characters,
verbose=True,
)
tokenizer.save(
args.tokenizer_output
)
print(
"\nTokenizer saved to:",
args.tokenizer_output,
)
print(
"Tokenizer fingerprint:",
tokenizer.fingerprint(),
)
loaded_tokenizer = (
BPETokenizer.from_file(
args.tokenizer_output
)
)
assert (
loaded_tokenizer.fingerprint()
== tokenizer.fingerprint()
)
print(
"Tokenizer reload test passed."
)
print(
"\nEncoding dataset splits..."
)
train_ids = validate_round_trip(
tokenizer=loaded_tokenizer,
text=train_text,
split_name="Training",
)
val_ids = validate_round_trip(
tokenizer=loaded_tokenizer,
text=val_text,
split_name="Validation",
)
test_ids = validate_round_trip(
tokenizer=loaded_tokenizer,
text=test_text,
split_name="Test",
)
output_dir = Path(
args.encoded_output_dir
)
train_tensor = save_token_tensor(
train_ids,
output_dir / "train_ids.pt",
)
val_tensor = save_token_tensor(
val_ids,
output_dir / "val_ids.pt",
)
test_tensor = save_token_tensor(
test_ids,
output_dir / "test_ids.pt",
)
for split_name, tensor in [
("Training", train_tensor),
("Validation", val_tensor),
("Test", test_tensor),
]:
assert tensor.min().item() >= 0
assert (
tensor.max().item()
< loaded_tokenizer.vocab_size
)
print(
f"{split_name} tensor: "
f"shape={tuple(tensor.shape)}, "
f"min={tensor.min().item()}, "
f"max={tensor.max().item()}"
)
character_tokens = len(
train_text
)
bpe_tokens = len(
train_ids
)
compression_factor = (
character_tokens / bpe_tokens
)
reduction_percentage = (
1.0
- bpe_tokens / character_tokens
) * 100.0
print(
"\nCharacter token count:",
f"{character_tokens:,}",
)
print(
"BPE token count:",
f"{bpe_tokens:,}",
)
print(
"Compression factor:",
f"{compression_factor:.2f}x",
)
print(
"Token-count reduction:",
f"{reduction_percentage:.2f}%",
)
learned_tokens = (
loaded_tokenizer.tokens[
len(
loaded_tokenizer.base_tokens
):
]
)
longest_tokens = sorted(
learned_tokens,
key=len,
reverse=True,
)
print("\nLongest learned tokens:")
for token in longest_tokens[:30]:
print(
f"{loaded_tokenizer.stoi5d} "
f"{token!r}"
)
sample = (
"Artificial intelligence "
"and machine learning"
)
unknown = [
char
for char in sample
if char not in (
loaded_tokenizer.stoi
)
]
if not unknown:
sample_tokens = (
loaded_tokenizer.encode_tokens(
sample
)
)
sample_ids = (
loaded_tokenizer.encode(
sample
)
)
print("\nSample text:")
print(sample)
print("\nSample BPE tokens:")
print(sample_tokens)
print("\nSample token IDs:")
print(sample_ids)
print("\nDecoded sample:")
print(
loaded_tokenizer.decode(
sample_ids
)
)
print(
"\nBPE tokenizer pipeline "
"completed successfully."
)
if __name__ == "__main__":
main()
Запуск
По умолчанию используется небольшой образовательный запуск:
python build_bpe_tokenizer.py
Чтобы изучить более обширный словарный запас:
python build_bpe_tokenizer.py \
--target-vocab-size 1000
Чтобы тренироваться на большем объёме текста также:
python build_bpe_tokenizer.py \
--training-characters 1000000 \
--target-vocab-size 1000
Для быстрой проверки, которая лишь подтверждает работоспособность конвейера:
python build_bpe_tokenizer.py \
--training-characters 100000 \
--target-vocab-size 300
Если ваша оболочка не поддерживает продолжение строк с использованием обратного слеша, введите каждую команду в отдельную строку.
Как выглядит вывод
Во время выполнения может появиться что-то вроде этого:
Training characters: 10,000,000
Validation characters: 1,100,000
Test characters: 1,200,000
BPE training characters: 500,000
Initial characters: 500,000
Base vocabulary size: 200
Merge 1: (' ', 't') → ' t', frequency=19,842
Merge 2: ('h', 'e') → 'he', frequency=17,415
Merge 3: (' t', 'he') → ' the', frequency=14,108
...
Final vocabulary size: 500
Первые объединения соответствуют тому, чего можно ожидать от английского языка: пробел за t, затем he, после чего они объединяются в " the". После кодирования резюме выглядит примерно так:
Character token count: 10,000,000
BPE token count: 4,500,000
Compression factor: 2.22x
Token-count reduction: 55.00%
Ваши показатели будут различаться. Обратите внимание, что обработка всего набора данных для обучения с использованием этого кодера может существенно увеличить время выполнения при больших словарях.
Что меняется в модели и её окне контекста
Количество токенов в реальном предложении
Возьмём это предложение:
The history of artificial intelligence
Токенизатор символов генерирует примерно столько токенов:
38 tokens
Токенизатор BPE может генерировать токены целых слов вроде этих:
["The", " history", " of", " artificial", " intelligence"]
Или же он может генерировать более мелкие фрагменты в зависимости от того, какие способы слияния он изучил. В любом случае количество будет находиться в следующем диапазоне:
5 to 15 tokens
Сегментация не создаётся вручную; она определяется на основе статистики обучающего корпуса.
Больше текста в окне контекста
Предположим, что модель использует такой размер блока:
block_size = 128
При использовании символов 128 токенов покрывают примерно 128 символов. С применением метода BPE те же 128 токенов могут охватывать несколько сотен символов. Модель получает больший языковой контекст без увеличения сложности вычислений внимания. Это важно, поскольку затраты на самовнимание растут примерно пропорционально:
T²
Здесь T — это длина последовательности; удвоение T приводит примерно к учетверению затрат на обработку вниманием, тогда как более эффективная токенизация позволяет вместить больше текста при той же длине T.
Компромисс между размером словаря и производительностью
Более большой словарь BPE обычно означает более короткие последовательности, но в то же время приводит к увеличению размера некоторых матриц. Таблица вкладок токенов содержит столько параметров:
vocab_size × embedding_dim
При связывании весов выходной головной модуль использует ту же матрицу, вместо того чтобы добавлять вторую. Сравните небольшой словарный запас:
Vocabulary: 500
Embedding dimension: 128
Embedding parameters: 64,000
С более крупным:
Vocabulary: 8,000
Embedding dimension: 128
Embedding parameters: 1,024,000
Таким образом, более большой словарь сокращает длину последовательностей, но увеличивает параметры эмбеддингов, стоимость вычислений логитов вывода, потребление памяти и время обучения токенизатора. Каждый дополнительный токен также требует достаточного количества вхождений для формирования качественного эмбеддинга, поэтому в корпусе размером с WikiText-2 огромный словарь приводит к появлению множества плохо обученных токенов. Оптимальный размер словаря зависит от набора данных и модели.
Повторное обучение модели с использованием токенов BPE
Обновите скрипт обучения, чтобы он загружал файл токенизатора:
tokenizer/bpe_tokenizer.json
И закодированные разделы:
data/encoded_bpe/train_ids.pt
data/encoded_bpe/val_ids.pt
data/encoded_bpe/test_ids.pt
Затем сформируйте конфигурацию на основе размера словаря токенизатора:
config = MiniGPTConfig(
vocab_size=bpe_tokenizer.vocab_size,
block_size=128,
embedding_dim=128,
num_heads=4,
num_layers=4,
expansion_factor=4,
dropout=0.1,
)
И сохраняйте точки контроля в свою собственную папку:
checkpoints/bpe/
Помимо этих изменений, цикл обучения может остаться практически прежним.
Справедливое сравнение двух моделей
Для значимого сравнения по возможности оставьте неизменными архитектуру (ширина, глубина, количество блоков обработки), данные, оптимизатор, график обучения, объем ресурсов, выделенных на обучение, и метод оценки. Даже в этом случае сравнение при одинаковом количестве шагов обучения вводит в заблуждение, поскольку пакет токенов BPE содержит значительно больше текста, чем пакет символов. Лучшими критериями для сравнения являются:
Total characters processed
Total bytes processed
Training time
Compute budget
Показатель перплексити вычисляется для каждой единицы прогнозирования, поэтому показатели перплексити для символов и BPE нельзя сравнивать напрямую. Чтобы получить сопоставимые значения, преобразуйте оба показателя в количество бит на символ; в противном случае оценивайте каждую модель по качеству генерации, рассчитывайте стоимость и метрики, соответствующие токенизатору.
Распространённые ошибки
Создание отдельного токенизатора для каждого разделения
Это неверно:
Train tokenizer from training text
Validation tokenizer from validation text
Test tokenizer from test text
Данные для обучения, проверки и тестирования должны быть закодированы с помощью одного и того же токенизатора; в противном случае один и тот же идентификатор будет означать разные вещи в разных разделах данных.
Применение объединений в неправильном порядке
Объединения строятся друг на друге, поэтому их необходимо выполнять ровно в том порядке, в котором они были обучены:
for pair in self.merges:
apply(pair)
Сортировка правил, например, в алфавитном порядке для лучшей читаемости, влияет на способ сегментации текста.
Запись файлов с кодировкой символов
Храните идентификаторы BPE в отдельном каталоге:
data/encoded_bpe/
И оставьте идентификаторы персонажей там, где они находятся:
data/encoded/
Обучение на основе текста для проверки
Статистика слияний должна поступать исключительно из учебного текста. Использование базового набора символов из каждой части текста допустимо для обеспечения полноты охвата, но частоты слияний должны отражать только учебный корпус; в противном случае показатели проверки будут завышенными.
Ожидание того, что выученные токены будут целыми словами
Метод BPE с удовольствием обучается на фрагментах вроде:
"ing"
"tion"
" th"
" arti"
Токены субслов сформированы на основе частоты встречаемости, а не по границам слов, установленным человеком, и это является ожидаемым поведением.
Предположение о том, что чем больше словарь, тем лучше
Большие словари увеличивают затраты на параметры и приводят к появлению редких токенов с недостаточной настройкой. Попробуйте несколько размеров словаря и измерьте результаты. Кроме того, никогда не сравнивайте валовую степень перплексности разных токенизаторов.
Упражнения
- Сравните три размера словаря. Обучайте модель с использованием указанных размеров и сравните время обучения, длину последовательностей, уровень сжатия, длину самых длинных токенов и размер векторов замещения.
300 tokens
500 tokens
1,000 tokens
- Изучите ранние правила слияния. Выведите первые 50 правил слияния. Встречаются ли шаблоны пробелов? Распространённые суффиксы? Комбинации знаков препинания? Какие правила слияния вас удивили?
- Токенизируйте одно предложение обоими способами. Закодируйте следующее предложение с помощью токенизатора на основе символов и токенизатора BPE, затем сравните полученные токены, их идентификаторы, количество токенов и результиры декодирования.
Artificial intelligence is transforming technology.
- Испытайте редкие слова. Закодируйте следующие слова и посмотрите, как BPE собирает неизвестные слова из известных фрагментов.
electromagnetism
neurocomputational
uncharacteristically
- Оцените охват контекста. Расшифруйте первые 128 токенов в каждом токенизаторе и определите, какой объем текста покрывает каждое окно.
- Обучите модель BPE. Обучайте её на файлах из приведённой ниже директории и сравните полученные примеры с моделью на основе символов при одинаковых вычислительных ресурсах.
data/encoded_bpe/
- Применение фингерпринтов. Хранить фингерпринт токенизатора в каждом чекпоинте и отказываться от выполнения инференса, если чекпоинт и действующий токенизатор не совпадают.
- Ускорение кодирования. Заменить кодер, выполняющий один проход при слиянии, на более быстрый вариант, использующий структуры типа мердж-рейнки, кучи, связанные списки токенов, кэши по словам или другие аналогичные методы.
Правила токенизации, которые стоит сохранить
- Субслова находятся между гибкими, но длинными символами и компактными, но хрупкими словами.
- Токенизатор BPE полностью определяется его базовыми токенами и порядком их слияний; необходимо сохранять именно эти элементы и создавать для них фингерпринты.
- При кодировании слияния выполняются в определенном порядке, при декодировании — токены соединяются, при этом процесс в обоих направлениях происходит без потерь.
- Один токенизатор используется для всех разбиений текста, а слияния обучаются исключительно на данных обучения.
Полная цепочка обработки, от сырого текста до переобученной модели, выглядит следующим образом:
WikiText training text
→ character vocabulary
→ pair frequencies
→ repeated merge operations
→ BPE vocabulary and merge rules
→ encode all splits
→ save token tensors
→ train a new Mini-GPT
Суть в том, что токенизация — это не просто деталь предварительной обработки. Она определяет, какие единицы будет изучать модель, сколько текста помещается в её окно контекста и какова сложность задачи генерации вывода. Как только небольшая версия GPT начинает работать с токенами BPE, естественным следующим шагом является изучение того, как та же самая система генерации следующих токенов развивается в современные системы генеративного ИИ:
Next-token predictor
→ pretrained language model
→ instruction-following model
→ retrieval-augmented application
→ tool-using GenAI system
Связанные статьи
- Отладка небольшого GPT в PyTorch: тесты, изолирующие каждую ошибку: пошаговая инструкция по отладке GPT на уровне символов в PyTorch, включающая работу с идентификаторами токенов и смещением целей, расчёт градиентов, обработку значений NaN и использование чекпоинтов.
- Создание ИИ-агента с нуля: шаблоны, ReAct и LangGraph: ознакомление с основными концепциями ИИ-агентов, планированием действий, использованием инструментов, рефлексией и шаблоном ReAct, а также с тем, как LangChain и LangGraph помогают вручную создать такой агент.