Кодирование и декодирование с помощью BPE-токенизатора

Предскажите первое слияние, прежде чем запускать программу

Возьмём пробел и следующую за ним кириллическую а:

текст             = " а"
байты UTF-8 (hex) = [20, d0, b0]
ID обучения       = [32, 208, 176]

Из главы 3 нужны два правила:

Какое правило сработает? Ранг задаёт приоритет; частоту пар для нового входа считать заново не нужно. Сначала правило ранга 0 заменяет пару 32,208, и остаётся [256,176]. После этой замены операнды правила ранга 1 уже не стоят рядом. Поскольку ID 0 и 1 зарезервированы для двух управляющих токенов, каждый ID из пространства обучения сдвигается на два:

канонические ID содержимого = [258, 178]
ID документа                = [0, 258, 178, 1]
                               ^             ^
                              BOS           EOS

Теперь восстановим содержимое. Для ID 258 сохранены байты 20 d0, а ID 178 соответствует исходному байту b0. Последовательное объединение даёт исходные байты 20 d0 b0. При этом последовательность 20 d0, сохранённая для первого токена, сама по себе не является корректным UTF-8: d0 начинает двухбайтовую кодировку, а байт продолжения b0 находится в следующем токене. Поэтому границы токенов не обязаны совпадать с границами символов.

Есть важная тонкость. [34,259] — тоже допустимая последовательность ID содержимого: ID 34 соответствует пробелу 20, а для ID 259, созданного слиянием ранга 1, сохранены байты d0 b0. Из неё восстанавливаются те же три байта. Но если снова закодировать эти байты, получится каноническая последовательность [258,178], потому что правило ранга 0 имеет приоритет. Декодирование без потерь ещё не делает любую допустимую последовательность ID канонической.

Главное: токенизатор гарантирует точное восстановление байтов для результатов собственного кодирования. Канонический результат кодирования и возможность декодировать произвольное допустимое разбиение — разные свойства.

Односторонняя гарантия: точное восстановление байтов

Токенизатор гарантирует:

encode_content сначала ставит в соответствие каждому байту один сдвинутый ID, а затем применяет зафиксированные правила по возрастанию ранга, начиная с нулевого. Функция не использует частоты из валидационной или тестовой выборки, не обучает новые пары, не нормализует Unicode, не меняет регистр и не делит строку на слова. decode_content последовательно объединяет байты, сохранённые для каждого токена содержимого.

Формула намеренно не утверждает, что $\operatorname{encode}_{content}(\operatorname{decode}_{content}(z))=z$ для любой последовательности $z$. Контрпример — [34,259]. Для обучения и последующего применения модели достаточно односторонней гарантии, записанной выше: кодировщик всегда выбирает одну и ту же каноническую последовательность, а декодировщик восстанавливает из неё исходные байты без потерь.

Обозначения и схема ID

В версии 1 используется единое непрерывное пространство ID:

В этом курсе ID для PAD не предусмотрен: в последующих главах используются окна фиксированной длины. Сценарии инференса, в которых последовательности приходится часто дополнять до общей длины (padding), оставлены за рамками курса; в рабочих системах токен PAD по-прежнему может быть необходим.

Как байтовый алфавит устраняет необходимость в <UNK>

Закрытый словарь целых слов может назначить отдельные ID частым словам из обучающих данных, но в нём нет записи для написания, которое раньше не встречалось. Исторический пример на Rust кодирует lowering, как и любое другое отсутствующее в словаре слово, идентификатором 0. Декодер возвращает лишь маркер <UNK> и уже не может восстановить исходное написание:

pub fn fit_closed_word_vocabulary(documents: &[&str]) -> BTreeMap<String, u32> {
    documents
        .iter()
        .flat_map(|document| document.split_whitespace())
        .collect::<BTreeSet<_>>()
        .into_iter()
        .enumerate()
        .map(|(index, word)| (word.to_owned(), index as u32 + 1))
        .collect()
}

/// Encodes a spelling, collapsing every unseen word to the same ID.
pub fn encode_closed_word(vocabulary: &BTreeMap<String, u32>, word: &str) -> u32 {
    vocabulary.get(word).copied().unwrap_or(UNKNOWN_WORD_ID)
}

/// Decodes a known ID or the literal unknown marker; original unseen bytes are gone.
pub fn decode_closed_word(vocabulary: &BTreeMap<String, u32>, token_id: u32) -> String {
    vocabulary
        .iter()
        .find_map(|(word, &id)| (id == token_id).then(|| word.clone()))
        .unwrap_or_else(|| "<UNK>".to_owned())
}

В работе 2016 года о редких словах Сеннрих, Хэддоу и Бёрч применили многократное слияние пар к последовательностям символов и получили подсловные единицы. Они занимают промежуточное положение между целыми словами и отдельными символами. Однако в этой работе ещё не предлагался универсальный байтовый алфавит.

В разделе о представлении входа отчёта GPT-2 описан следующий шаг: базовый алфавит из 256 значений — по одному на каждый возможный байт — позволяет представить любую строку Unicode в кодировке UTF-8. За такую универсальность приходится платить более длинными последовательностями. До слияний кириллическая буква занимает два байта UTF-8, а эмодзи — часто четыре. Обученные слияния сокращают часто встречающиеся последовательности, но байтовая основа не гарантирует ни минимального итогового словаря, ни одинаковой длины токенизации для разных письменностей.

В GPT-2 перед применением BPE текст дополнительно разбивался по категориям символов, а пробел обрабатывался особым образом. Наш токенизатор устроен иначе: он применяет конкретную таблицу рангов из главы 3, разрешает слияние любой пары внутри документа и никогда не пересекает его границу. Поэтому название «BPE» само по себе ещё не определяет все важные правила реализации.

Исторический вывод: байтовая основа устраняет проблему неизвестных строк, а конкретное поведение токенизатора по-прежнему определяется приоритетом слияний и правилами границ.

Реализуйте на Rust зафиксированный токенизатор со строгой проверкой

Схема ID закрепляет разделение управляющих токенов и содержимого в самом коде, а не только в комментариях. Реализация проверяет, что последний ID помещается в u32, и предоставляет диапазоны ID для байтов и результатов слияний:

/// Serialized layout version taught by Chapter 4.
pub const TOKENIZER_LAYOUT_VERSION: u32 = 1;
/// Marks the beginning of one encoded document.
pub const BOS_TOKEN_ID: u32 = 0;
/// Marks the end of one encoded document.
pub const EOS_TOKEN_ID: u32 = 1;
/// Maps every Chapter 3 training-space ID into the content namespace.
pub const CONTENT_ID_OFFSET: u32 = 2;
/// First content ID representing one raw byte.
pub const FIRST_BYTE_TOKEN_ID: u32 = CONTENT_ID_OFFSET;
/// Last content ID representing one raw byte.
pub const LAST_BYTE_TOKEN_ID: u32 = CONTENT_ID_OFFSET + BYTE_TOKEN_COUNT - 1;
/// Content ID assigned to merge rank zero.
pub const FIRST_MERGE_TOKEN_ID: u32 = LAST_BYTE_TOKEN_ID + 1;

/// The fixed fields and vocabulary extent of tokenizer layout version 1.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TokenizerLayout {
    merge_count: usize,
    vocabulary_size: usize,
}

impl TokenizerLayout {
    /// Validates that every content symbol can be represented by a `u32` ID.
    pub fn new(merge_count: usize) -> Result<Self, BpeTokenizerError> {
        let vocabulary_size = usize::try_from(FIRST_MERGE_TOKEN_ID)
            .ok()
            .and_then(|base| base.checked_add(merge_count))
            .ok_or(BpeTokenizerError::LayoutOverflow { merge_count })?;
        let highest_token = vocabulary_size
            .checked_sub(1)
            .and_then(|value| u32::try_from(value).ok())
            .ok_or(BpeTokenizerError::LayoutOverflow { merge_count })?;
        if merge_count > 0 && highest_token < FIRST_MERGE_TOKEN_ID {
            return Err(BpeTokenizerError::LayoutOverflow { merge_count });
        }
        Ok(Self {
            merge_count,
            vocabulary_size,
        })
    }

    /// Returns the serialized layout version.
    pub const fn version(self) -> u32 {
        TOKENIZER_LAYOUT_VERSION
    }

    /// Returns the number of frozen merge ranks.
    pub const fn merge_count(self) -> usize {
        self.merge_count
    }

    /// Returns the complete number of control and content IDs.
    pub const fn vocabulary_size(self) -> usize {
        self.vocabulary_size
    }

    /// Maps a raw byte to its one-byte content token.
    pub const fn byte_token_id(self, byte: u8) -> u32 {
        FIRST_BYTE_TOKEN_ID + byte as u32
    }

    /// Returns the final content ID assigned to a valid zero-based rank.
    pub fn merge_token_id(self, rank: usize) -> Option<u32> {
        if rank >= self.merge_count {
            return None;
        }
        usize::try_from(FIRST_MERGE_TOKEN_ID)
            .ok()
            .and_then(|base| base.checked_add(rank))
            .and_then(|token| u32::try_from(token).ok())
    }
}

BpeTraining нельзя получить из произвольной таблицы пар через публичный API. На каждом шаге BpeTrainer::train создаёт согласованный набор данных: ранг правила, пару операндов, ID нового токена в пространстве обучения и байтовое представление этого токена. Поля BpeTraining недоступны вызывающему коду напрямую, а публичные методы предоставляют эти данные только для чтения. Поэтому после успешного обучения вызывающий код не может подменить или изменить их.

BpeTokenizer::from_training доверяет этим уже проверенным данным. Метод копирует ранг, пару и ID нового токена каждого правила, а также готовый словарь байтовых представлений. При копировании он переводит ID операндов и результата в схему версии 1, прибавляя два. Метод не строит байтовые представления заново и не проверяет повторно инварианты из главы 3. Он проверяет только новое условие главы 4: после резервирования BOS и EOS наибольший ID содержимого, вычисленный по числу полученных при обучении правил, должен помещаться в диапазон u32. Токенизатор владеет скопированными данными, поэтому продолжает работать после того, как исходное значение BpeTraining перестанет существовать.

BpeTokenizer::from_merge_pairs, напротив, получает таблицу пар напрямую от вызывающего кода, поэтому проверяет все её инварианты. Сначала метод вычисляет наибольший ID, необходимый для схемы версии 1, и проверяет, что он помещается в диапазон u32. На шаге ранга $r$ метод назначает новому токену ID $256+r$ в пространстве обучения. Оба операнда должны быть определены до этого шага, а одна и та же пара не может встречаться на двух разных рангах. Байтовое представление нового токена строится объединением байтовых представлений операндов. После этого ID обоих операндов и результата увеличиваются на два и сохраняются в пространстве ID содержимого. Упражнения и код загрузки контрольных точек вызывают этот конструктор, когда у них есть таблица пар, но нет готового значения BpeTraining.

Основной метод кодировщика принимает &[u8]; вариант для &str передаёт ему байты UTF-8. К моменту кодирования каждое правило уже записано в пространстве ID содержимого: ID обоих операндов и результата увеличены на $2$. Кодировщик применяет эти готовые правила по возрастанию ранга; каждое правило выполняет один проход слева направо, как в главе 3.

Обычный метод и метод с трассировкой используют один и тот же цикл: он перебирает правила в порядке рангов и выполняет слияния. Обычный метод не сохраняет промежуточные состояния и возвращает только итоговые ID токенов. Метод с трассировкой дополнительно копирует последовательность до и после каждого правила, которое изменило входные данные. Эти снимки позволяют изучить ход кодирования, но не влияют на результат:

    fn initial_content_tokens(&self, bytes: &[u8]) -> Vec<u32> {
        bytes
            .iter()
            .map(|byte| self.layout.byte_token_id(*byte))
            .collect()
    }

    fn apply_ranked_merges(
        &self,
        mut content_tokens: Vec<u32>,
        mut observe: impl FnMut(&BpeMergeRule, usize, &[u32], &[u32]),
    ) -> Vec<u32> {
        for rule in &self.merge_rules {
            let before = content_tokens;
            let (after, replacements) =
                replace_pair_left_to_right(&before, rule.content_pair, rule.content_token_id);
            if replacements > 0 {
                observe(rule, replacements, &before, &after);
            }
            content_tokens = after;
        }
        content_tokens
    }

    /// Encodes bytes and records every rank that changed the sequence.
    pub fn encode_content_with_trace(&self, bytes: &[u8]) -> BpeEncodingTrace {
        let initial_tokens = self.initial_content_tokens(bytes);
        let mut applications = Vec::new();
        let content_tokens = self.apply_ranked_merges(
            initial_tokens.clone(),
            |rule, replacements, before, after| {
                applications.push(BpeMergeApplication {
                    rank: rule.rank,
                    replacements,
                    before: before.to_vec(),
                    after: after.to_vec(),
                });
            },
        );

        BpeEncodingTrace {
            initial_tokens,
            applications,
            content_tokens,
        }
    }

    /// Encodes arbitrary bytes into the canonical rank-ordered content sequence.
    pub fn encode_content(&self, bytes: &[u8]) -> Vec<u32> {
        let initial_tokens = self.initial_content_tokens(bytes);
        self.apply_ranked_merges(initial_tokens, |_, _, _, _| {})
    }

    /// Encodes a valid UTF-8 string through the same byte boundary.
    pub fn encode_utf8(&self, text: &str) -> Vec<u32> {
        self.encode_content(text.as_bytes())
    }

Основной метод декодирования возвращает Vec<u8>. Благодаря этому можно без потерь восстановить любое из 256 значений байта. Например, последовательность ff fe восстанавливается как байты, хотя она не является корректным UTF-8. Отдельный метод преобразует результат в String и, если UTF-8 некорректен, возвращает позицию ошибочного байта вместо молчаливой подстановки символа:

    /// Concatenates content-token expansions without interpreting them as text.
    pub fn decode_content(&self, content: &[u32]) -> Result<Vec<u8>, BpeTokenizerError> {
        for (position, &token_id) in content.iter().enumerate() {
            if token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID {
                return Err(BpeTokenizerError::ControlTokenInContent { position, token_id });
            }
        }
        self.decode_tokens(content, 0)
    }

    /// Decodes content bytes, then requires the result to be valid UTF-8.
    pub fn decode_content_utf8(&self, content: &[u32]) -> Result<String, BpeTokenizerError> {
        strict_utf8(self.decode_content(content)?)
    }

    /// Decodes a wrapped document, then requires the result to be valid UTF-8.
    pub fn decode_document_utf8(&self, document: &[u32]) -> Result<String, BpeTokenizerError> {
        strict_utf8(self.decode_document(document)?)
    }

Управляющие токены задают структуру документа отдельно от его содержимого. Кодировщик сначала выполняет все слияния и лишь затем добавляет [BOS, ..., EOS], поэтому управляющий ID не может стать операндом правила. Декодер требует обе границы и отклоняет BOS или EOS во внутренней позиции:

    /// Encodes content first, then adds controls that never enter a merge pass.
    pub fn encode_document(&self, bytes: &[u8]) -> Vec<u32> {
        let content = self.encode_content(bytes);
        let mut document = Vec::with_capacity(content.len() + 2);
        document.push(BOS_TOKEN_ID);
        document.extend(content);
        document.push(EOS_TOKEN_ID);
        document
    }

    /// Encodes a UTF-8 document and adds its endpoint controls.
    pub fn encode_utf8_document(&self, text: &str) -> Vec<u32> {
        self.encode_document(text.as_bytes())
    }

    /// Validates one wrapped document and recovers its exact content bytes.
    pub fn decode_document(&self, document: &[u32]) -> Result<Vec<u8>, BpeTokenizerError> {
        if document.len() < 2 {
            return Err(BpeTokenizerError::DocumentTooShort {
                length: document.len(),
            });
        }
        if document[0] != BOS_TOKEN_ID {
            return Err(BpeTokenizerError::ExpectedBos { found: document[0] });
        }
        let last = document.len() - 1;
        if document[last] != EOS_TOKEN_ID {
            return Err(BpeTokenizerError::ExpectedEos {
                found: document[last],
            });
        }
        for (position, &token_id) in document[1..last].iter().enumerate() {
            if token_id == BOS_TOKEN_ID || token_id == EOS_TOKEN_ID {
                return Err(BpeTokenizerError::InteriorControlToken {
                    position: position + 1,
                    token_id,
                });
            }
        }
        self.decode_tokens(&document[1..last], 1)
    }

Демонстрационная программа ch04-apply-bpe-tokenizer, исходный код которой показан ниже, заново обучает BPE на обучающей выборке, получает те же восемь ранжированных правил, что и в главе 3, а затем фиксирует их в токенизаторе. Программа показывает потерю исходного написания в историческом словаре целых слов, кодирует строку 🦀, не встречавшуюся при обучении, без потерь восстанавливает последовательность байтов ff fe, которая не является корректным UTF-8, и демонстрирует канонизацию. Она выводит схему ID токенизатора, результаты проверок граничных случаев и два примера точного восстановления исходных байтов:

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let corpus = Corpus::from_json(CORPUS_JSON)?;
    let manifest = SplitManifest::from_json(SPLIT_MANIFEST)?;
    let partitions = manifest.partition(&corpus)?;
    let training = BpeTrainer::new(8).train(&partitions)?;
    let tokenizer = BpeTokenizer::from_training(&training)?;
    let layout = tokenizer.layout();

    println!("layout version: {}", layout.version());
    println!("control ids: BOS={BOS_TOKEN_ID} EOS={EOS_TOKEN_ID}");
    println!(
        "content ranges: bytes={FIRST_BYTE_TOKEN_ID}..{LAST_BYTE_TOKEN_ID} merges={FIRST_MERGE_TOKEN_ID}..{} vocabulary={}",
        layout.vocabulary_size() - 1,
        layout.vocabulary_size()
    );
    println!("frozen merge ranks: {} (train only)", layout.merge_count());

    let historical = fit_closed_word_vocabulary(&["low lower", "new newest"]);
    let historical_id = encode_closed_word(&historical, "lowering");
    println!("historical input: lowering");
    println!("historical ids: {}", format_ids(&[historical_id]));
    println!(
        "historical decoded: {} (original bytes lost)",
        decode_closed_word(&historical, historical_id)
    );

    let unseen = "🦀";
    let unseen_ids = tokenizer.encode_utf8(unseen);
    println!("unseen UTF-8 input: {unseen}");
    println!("unseen content ids: {}", format_ids(&unseen_ids));
    println!(
        "unseen decoded: {}",
        tokenizer.decode_content_utf8(&unseen_ids)?
    );
    println!(
        "empty document ids: {}",
        format_ids(&tokenizer.encode_document(b""))
    );

    let malformed = [0xff, 0xfe];
    let malformed_ids = tokenizer.encode_content(&malformed);
    println!("malformed input bytes: {}", format_hex_array(&malformed));
    println!("malformed content ids: {}", format_ids(&malformed_ids));
    println!(
        "malformed decoded bytes: {}",
        format_hex_array(&tokenizer.decode_content(&malformed_ids)?)
    );
    println!(
        "malformed UTF-8: {}",
        tokenizer
            .decode_content_utf8(&malformed_ids)
            .expect_err("malformed bytes are not text")
    );
    println!(
        "interior control rejected: {}",
        tokenizer
            .decode_document(&[BOS_TOKEN_ID, 100, BOS_TOKEN_ID, EOS_TOKEN_ID])
            .expect_err("interior BOS must fail")
    );

    let canonical = tokenizer.encode_utf8(" а");
    let noncanonical = [34, 259];
    let same_bytes = tokenizer.decode_content(&noncanonical)?;
    println!("canonical \" а\" content ids: {}", format_ids(&canonical));
    println!("noncanonical same-byte ids: {}", format_ids(&noncanonical));
    println!(
        "noncanonical re-encodes as: {}",
        format_ids(&tokenizer.encode_content(&same_bytes))
    );

    println!("TRACE apply-bpe-tokenizer-v1 BEGIN");
    println!(
        "LAYOUT version={} bos={} eos={} content_offset=2 byte_count=256 merge_count={} vocabulary_size={}",
        layout.version(),
        BOS_TOKEN_ID,
        EOS_TOKEN_ID,
        layout.merge_count(),
        layout.vocabulary_size()
    );
    for rule in tokenizer.merge_rules() {
        println!(
            "RULE rank={} training_pair={},{} training_token={} content_pair={},{} content_token={} bytes_hex={}",
            rule.rank(),
            rule.training_pair().left(),
            rule.training_pair().right(),
            rule.training_token_id(),
            rule.content_pair().left(),
            rule.content_pair().right(),
            rule.content_token_id(),
            format_hex(
                tokenizer
                    .token_bytes(rule.content_token_id())
                    .expect("merge token has bytes")
            )
        );
    }
    print_trace_case(&tokenizer, "ascii-bee", b"bee ")?;
    print_trace_case(&tokenizer, "cyrillic-a", " а".as_bytes())?;
    println!("TRACE apply-bpe-tokenizer-v1 END");
    println!("chapter 5 handoff: preserve each wrapped document boundary");
    Ok(())
}

BpeTokenizer также обрабатывает пустое содержимое, все 256 значений байта, ASCII, кириллицу, эмодзи, сочетания с NUL и переводом строки, а также некорректный UTF-8. Для каждого вида недопустимых данных предусмотрен отдельный вариант BpeTokenizerError: переполнение пространства ID, ссылка правила слияния на ещё не определённый операнд, повтор пары, неизвестный ID токена или неверное положение управляющего токена. Эти граничные случаи важны: гарантия точного восстановления относится и к данным, которые не являются корректным текстом, однако структура документа всё равно должна запрещать управляющие токены в неверных позициях.

Проследите группировку и восстановление исходных байтов

Оба примера показывают результаты, вычисленные программой на Rust. В строке bee слияние ранга 7 объединяет последнюю e с пробелом в токен 265, а предыдущие b и e остаются однобайтовыми токенами. В строке а слияние ранга 0 объединяет пробел с первым байтом буквы а; байт продолжения остаётся отдельно. Второй пример наглядно показывает, что границы токенов и символов могут не совпадать.

Сначала пройдите по каждой схеме сверху вниз и проследите кодирование. Затем для декодирования объедините сохранённые байты токенов слева направо, не меняя их порядок. BOS и EOS не входят в содержимое. Равенство исходных и восстановленных байтов обозначено не только цветом, но и галочкой, рамкой, подписью и самими значениями байтов. В обоих примерах объединение декодированных байтов содержимого в точности восстанавливает исходную последовательность байтов.

Сначала предскажите, затем проверьте

1. Для байтов 20 d0 b0 предскажите результат при ошибочном применении ранга 1 до ранга 0, а затем запишите канонический результат.
2. Переведите ASCII ? (байт 3f), ID слияния 257 из пространства обучения и ранг 7 в схему версии 1.
3. Предскажите encode_content("") и encode_document("").
4. Предскажите ID содержимого для 🦀, если ни одна обученная пара не совпадает.
5. Кириллическая т кодируется в UTF-8 как d1 82, а правило ранга 2 объединяет пару (209,130). Предскажите конечную последовательность.
6. Декодируйте [34,259], затем предскажите результат повторного кодирования.
7. Найдите первую ошибку в документе [0,100,0,1] и содержимом [100,1].
8. Решите, обязано ли точное восстановление байтов из ID [257,256] давать корректный String.

Распространённое заблуждение: «Если последовательность декодируется без потерь, то повторное кодирование восстановленных байтов обязательно вернёт ту же последовательность токенов». Это неверно. Несколько вариантов разбиения могут декодироваться без потерь, но зафиксированный порядок рангов выбирает только одно каноническое кодирование.

Сохраняйте каждый документ как отдельную последовательность с границами

Теперь общая реализация курса преобразует каждый исходный документ в отдельную воспроизводимую последовательность [BOS, содержимое..., EOS], из которой можно восстановить исходные байты. Статистика слияний по-прежнему вычисляется только по обучающим документам из главы 2. Применение зафиксированного токенизатора к валидационной или тестовой выборке не изменяет его словарь и ранги.

Глава 5 будет обрабатывать эти последовательности по одной. Пары «вход — цель» будут строиться только внутри каждого документа, без объединения документов и пересечения границ между обучающей, валидационной и тестовой выборками. BOS задаёт начальный контекст, а EOS становится последней целью предсказания. Оба токена явно обозначают границы документа и не участвуют в слияниях.

Переход к следующей главе: теперь глава 5 может строить обучающие пары из полных обратимых последовательностей с явными границами документов, не меняя байты их содержимого.