15.9. Набор данных для предобучения BERT

Для предобучения модели BERT, реализованной в Section 15.8, нам необходимо сформировать набор данных в оптимальном формате, чтобы обеспечить выполнение двух задач предобучения: маскированного языкового моделирования и предсказания следующего предложения. С одной стороны, оригинальная модель BERT предобучалась на объединении двух огромных корпусов — BookCorpus и английской Wikipedia (см. Section 15.8.5), что делает её запуск затруднительным для большинства читателей этой книги. С другой стороны, готовая предобученная модель BERT может не подходить для приложений из специфических областей, таких как медицина. Поэтому предобучение BERT на собственном наборе данных становится всё более популярным. Чтобы продемонстрировать процесс предобучения BERT, мы используем более компактный корпус WikiText-2 (Merity et al., 2016).

В сравнении с набором данных PTB, использованным для предобучения word2vec в Section 15.3, WikiText-2 (i) сохраняет исходную пунктуацию, что делает его подходящим для предсказания следующего предложения; (ii) сохраняет исходный регистр и числа; (iii) более чем в два раза больше.

import os
import random
import torch
from d2l import torch as d2l

В наборе данных WikiText-2 каждая строка представляет собой абзац, в котором пробелы вставляются между любыми знаками препинания и предшествующими им токенами. Сохраняются только абзацы, содержащие как минимум два предложения. Для простоты мы используем только точку в качестве разделителя при разбиении на предложения. Обсуждение более сложных методов разбиения на предложения оставлено для упражнений в конце этого раздела.

#@save
d2l.DATA_HUB['wikitext-2'] = (
    'https://s3.amazonaws.com/research.metamind.io/wikitext/'
    'wikitext-2-v1.zip', '3c914d17d80b1459be871a5039ac23e752a53cbe')

#@save
def _read_wiki(data_dir):
    file_name = os.path.join(data_dir, 'wiki.train.tokens')
    with open(file_name, 'r') as f:
        lines = f.readlines()
    # Uppercase letters are converted to lowercase ones
    paragraphs = [line.strip().lower().split(' . ')
                  for line in lines if len(line.split(' . ')) >= 2]
    random.shuffle(paragraphs)
    return paragraphs

15.9.1. Определение вспомогательных функций для задач предобучения

В дальнейшем мы начнем с реализации вспомогательных функций для двух задач предобучения BERT: предсказания следующего предложения и маскированного языкового моделирования. Эти вспомогательные функции будут вызываться позже при преобразовании необработанного текстового корпуса в набор данных идеального формата для предобучения BERT.

15.9.1.1. Формирование задачи предсказания следующего предложения

Согласно описаниям Section 15.8.5.2, функция _get_next_sentence генерирует пример для задачи бинарной классификации.

#@save
def _get_next_sentence(sentence, next_sentence, paragraphs):
    if random.random() < 0.5:
        is_next = True
    else:
        # `paragraphs` is a list of lists of lists
        next_sentence = random.choice(random.choice(paragraphs))
        is_next = False
    return sentence, next_sentence, is_next

Следующая функция генерирует обучающие примеры для задачи предсказания следующего предложения из входных данных paragraph, вызывая функцию _get_next_sentence. Здесь paragraph — это список предложений, где каждое предложение является списком токенов. Аргумент max_len задаёт максимальную длину последовательности входных данных BERT в процессе предобучения.

#@save
def _get_nsp_data_from_paragraph(paragraph, paragraphs, vocab, max_len):
    nsp_data_from_paragraph = []
    for i in range(len(paragraph) - 1):
        tokens_a, tokens_b, is_next = _get_next_sentence(
            paragraph[i], paragraph[i + 1], paragraphs)
        # Consider 1 '<cls>' token and 2 '<sep>' tokens
        if len(tokens_a) + len(tokens_b) + 3 > max_len:
            continue
        tokens, segments = d2l.get_tokens_and_segments(tokens_a, tokens_b)
        nsp_data_from_paragraph.append((tokens, segments, is_next))
    return nsp_data_from_paragraph

15.9.1.2. Генерация задачи маскированного языкового моделирования

Для генерации обучающих примеров для задачи маскированного языкового моделирования из входной последовательности BERT мы определяем следующую _replace_mlm_tokens функцию. На вход tokens подается список токенов, представляющий входную последовательность BERT, candidate_pred_positions — список индексов токенов входной последовательности BERT, исключающий индексы специальных токенов (специальные токены не предсказываются в задаче маскированного языкового моделирования), а num_mlm_preds указывает количество предсказаний (напомним, что предсказываются 15% случайных токенов). В соответствии с определением задачи маскированного языкового моделирования в Section 15.8.5.1, на каждой позиции предсказания входной токен может быть заменен специальным токеном «<mask>» или случайным токеном, либо остаться без изменений. В результате функция возвращает входные токены после возможной замены, индексы токенов, для которых выполняются предсказания, и метки для этих предсказаний.

#@save
def _replace_mlm_tokens(tokens, candidate_pred_positions, num_mlm_preds,
                        vocab):
    # For the input of a masked language model, make a new copy of tokens and
    # replace some of them by '<mask>' or random tokens
    mlm_input_tokens = [token for token in tokens]
    pred_positions_and_labels = []
    # Shuffle for getting 15% random tokens for prediction in the masked
    # language modeling task
    random.shuffle(candidate_pred_positions)
    for mlm_pred_position in candidate_pred_positions:
        if len(pred_positions_and_labels) >= num_mlm_preds:
            break
        masked_token = None
        # 80% of the time: replace the word with the '<mask>' token
        if random.random() < 0.8:
            masked_token = '<mask>'
        else:
            # 10% of the time: keep the word unchanged
            if random.random() < 0.5:
                masked_token = tokens[mlm_pred_position]
            # 10% of the time: replace the word with a random word
            else:
                masked_token = random.choice(vocab.idx_to_token)
        mlm_input_tokens[mlm_pred_position] = masked_token
        pred_positions_and_labels.append(
            (mlm_pred_position, tokens[mlm_pred_position]))
    return mlm_input_tokens, pred_positions_and_labels

Вызывая описанную выше _replace_mlm_tokens функцию, следующая функция принимает на вход последовательность BERT (tokens) и возвращает индексы входных токенов (после возможной замены токенов, описанной в Section 15.8.5.1), индексы токенов, для которых выполняются предсказания, и индексы меток для этих предсказаний.

#@save
def _get_mlm_data_from_tokens(tokens, vocab):
    candidate_pred_positions = []
    # `tokens` is a list of strings
    for i, token in enumerate(tokens):
        # Special tokens are not predicted in the masked language modeling
        # task
        if token in ['<cls>', '<sep>']:
            continue
        candidate_pred_positions.append(i)
    # 15% of random tokens are predicted in the masked language modeling task
    num_mlm_preds = max(1, round(len(tokens) * 0.15))
    mlm_input_tokens, pred_positions_and_labels = _replace_mlm_tokens(
        tokens, candidate_pred_positions, num_mlm_preds, vocab)
    pred_positions_and_labels = sorted(pred_positions_and_labels,
                                       key=lambda x: x[0])
    pred_positions = [v[0] for v in pred_positions_and_labels]
    mlm_pred_labels = [v[1] for v in pred_positions_and_labels]
    return vocab[mlm_input_tokens], pred_positions, vocab[mlm_pred_labels]

15.9.2. Преобразование текста в набор данных для предобучения

Теперь мы почти готовы создать класс Dataset для предобучения BERT. Перед этим нам всё ещё нужно определить вспомогательную функцию _pad_bert_inputs для добавления специальных токенов «<pad>» к входным данным. Её аргумент examples содержит результаты работы вспомогательных функций _get_nsp_data_from_paragraph и _get_mlm_data_from_tokens для двух задач предобучения.

#@save
def _pad_bert_inputs(examples, max_len, vocab):
    max_num_mlm_preds = round(max_len * 0.15)
    all_token_ids, all_segments, valid_lens,  = [], [], []
    all_pred_positions, all_mlm_weights, all_mlm_labels = [], [], []
    nsp_labels = []
    for (token_ids, pred_positions, mlm_pred_label_ids, segments,
         is_next) in examples:
        all_token_ids.append(torch.tensor(token_ids + [vocab['<pad>']] * (
            max_len - len(token_ids)), dtype=torch.long))
        all_segments.append(torch.tensor(segments + [0] * (
            max_len - len(segments)), dtype=torch.long))
        # `valid_lens` excludes count of '<pad>' tokens
        valid_lens.append(torch.tensor(len(token_ids), dtype=torch.float32))
        all_pred_positions.append(torch.tensor(pred_positions + [0] * (
            max_num_mlm_preds - len(pred_positions)), dtype=torch.long))
        # Predictions of padded tokens will be filtered out in the loss via
        # multiplication of 0 weights
        all_mlm_weights.append(
            torch.tensor([1.0] * len(mlm_pred_label_ids) + [0.0] * (
                max_num_mlm_preds - len(pred_positions)),
                dtype=torch.float32))
        all_mlm_labels.append(torch.tensor(mlm_pred_label_ids + [0] * (
            max_num_mlm_preds - len(mlm_pred_label_ids)), dtype=torch.long))
        nsp_labels.append(torch.tensor(is_next, dtype=torch.long))
    return (all_token_ids, all_segments, valid_lens, all_pred_positions,
            all_mlm_weights, all_mlm_labels, nsp_labels)

Объединив вспомогательные функции для генерации обучающих примеров двух задач предобучения и вспомогательную функцию для дополнения входных данных, мы определяем следующий _WikiTextDataset класс как набор данных WikiText-2 для предобучения BERT. Реализовав функцию __getitem__, мы можем произвольным образом обращаться к примерам предобучения (маскированного языкового моделирования и предсказания следующего предложения), сгенерированным из пары предложений из корпуса WikiText-2.

В оригинальной модели BERT используются эмбеддинги WordPiece, размер словаря которых составляет 30000 (Wu et al., 2016). Метод токенизации WordPiece представляет собой незначительную модификацию исходного алгоритма кодирования пар байтов, описанного в Section 15.6.2. Для простоты мы используем функцию d2l.tokenize для токенизации. Токены, встречающиеся реже пяти раз, отфильтровываются.

#@save
class _WikiTextDataset(torch.utils.data.Dataset):
    def __init__(self, paragraphs, max_len):
        # Input `paragraphs[i]` is a list of sentence strings representing a
        # paragraph; while output `paragraphs[i]` is a list of sentences
        # representing a paragraph, where each sentence is a list of tokens
        paragraphs = [d2l.tokenize(
            paragraph, token='word') for paragraph in paragraphs]
        sentences = [sentence for paragraph in paragraphs
                     for sentence in paragraph]
        self.vocab = d2l.Vocab(sentences, min_freq=5, reserved_tokens=[
            '<pad>', '<mask>', '<cls>', '<sep>'])
        # Get data for the next sentence prediction task
        examples = []
        for paragraph in paragraphs:
            examples.extend(_get_nsp_data_from_paragraph(
                paragraph, paragraphs, self.vocab, max_len))
        # Get data for the masked language model task
        examples = [(_get_mlm_data_from_tokens(tokens, self.vocab)
                      + (segments, is_next))
                     for tokens, segments, is_next in examples]
        # Pad inputs
        (self.all_token_ids, self.all_segments, self.valid_lens,
         self.all_pred_positions, self.all_mlm_weights,
         self.all_mlm_labels, self.nsp_labels) = _pad_bert_inputs(
            examples, max_len, self.vocab)

    def __getitem__(self, idx):
        return (self.all_token_ids[idx], self.all_segments[idx],
                self.valid_lens[idx], self.all_pred_positions[idx],
                self.all_mlm_weights[idx], self.all_mlm_labels[idx],
                self.nsp_labels[idx])

    def __len__(self):
        return len(self.all_token_ids)

Используя функцию _read_wiki и класс _WikiTextDataset, мы определяем следующую функцию load_data_wiki для скачивания набора данных WikiText-2 и генерации примеров для предобучения из него.

#@save
def load_data_wiki(batch_size, max_len):
    """Load the WikiText-2 dataset."""
    num_workers = d2l.get_dataloader_workers()
    data_dir = d2l.download_extract('wikitext-2', 'wikitext-2')
    paragraphs = _read_wiki(data_dir)
    train_set = _WikiTextDataset(paragraphs, max_len)
    train_iter = torch.utils.data.DataLoader(train_set, batch_size,
                                        shuffle=True, num_workers=num_workers)
    return train_iter, train_set.vocab

Задавая размер батча равным 512 и максимальную длину входной последовательности BERT равной 64, мы выводим формы мини-батча примеров предобучения BERT. Обратите внимание, что в каждой входной последовательности BERT для задачи маскированного языкового моделирования предсказываются позиции \(10\) (\(64 \times 0.15\)).

batch_size, max_len = 512, 64
train_iter, vocab = load_data_wiki(batch_size, max_len)

for (tokens_X, segments_X, valid_lens_x, pred_positions_X, mlm_weights_X,
     mlm_Y, nsp_y) in train_iter:
    print(tokens_X.shape, segments_X.shape, valid_lens_x.shape,
          pred_positions_X.shape, mlm_weights_X.shape, mlm_Y.shape,
          nsp_y.shape)
    break

В заключение рассмотрим размер словаря. Даже после фильтрации редко встречающихся токенов он всё ещё более чем в два раза превышает размер словаря набора данных PTB.

len(vocab)

15.9.3. Итоги

• В сравнении с набором данных PTB набор данных WikiText-2 сохраняет исходную пунктуацию, регистр и цифры, а также более чем в два раза больше по объёму.

• Мы можем произвольно обращаться к примерам предобучения (маскированного языкового моделирования и предсказания следующего предложения), сгенерированным из пары предложений из корпуса WikiText-2.

15.9.4. Упражнения

1. Для простоты в качестве единственного разделителя для разбиения предложений используется точка. Попробуйте другие методы разбиения предложений, например, с помощью spaCy и NLTK. Возьмем NLTK в качестве примера. Сначала необходимо установить NLTK: pip install nltk. В коде сначала import nltk. Затем скачайте токенизатор предложений Punkt: nltk.download('punkt'). Чтобы разбить предложение, такое как sentences = 'This is great ! Why not ?', вызов функции nltk.tokenize.sent_tokenize(sentences) вернет список из двух строк с предложениями: ['This is great !', 'Why not ?'].

2. Каков размер словаря, если не отфильтровывать редкие токены?

Обсуждения