pretraining_data.py

# Copyright (c) 2021 Graphcore Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import glob
import multiprocessing
import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import numpy as np
import torch
from torch.utils.data import IterableDataset, Dataset
from poptorch import DataLoader
from poptorch.enums import DataLoaderMode
import popdist
from transformers import BertTokenizerFast
from tfrecord.reader import tfrecord_loader


TFRECORD_KEYS = (  # Torch Model Keys
    "input_ids",  # input_ids                  : tokens after masking
    "input_mask",  # attention_mask             : 0 if padded token, 1 otherwise
    "segment_ids",  # token_type_ids             : sentence 0 or 1
    "masked_lm_positions",  # masked_lm_positions        : position of masked tokens in input_ids
    "masked_lm_ids",  # masked_lm_labels=None      : label of masked tokens with padding as 0
    "next_sentence_labels",  # next_sentence_label=None   : 1 if next sentence, 0 otherwise
)

TFRECORD_KEYS_PACKED = (
    "packed_input_ids",  # : tokens after masking
    "packed_input_mask",  # : 0 if padded token. 1, 2 or 3 if token belongs to 1st, 2nd or 3rd sequence resp.
    "packed_segment_ids",  # : sentence 0 or 1 for each sequence in the pack
    "packed_position_ids",  # : position of tokens relative to each sequence
    "packed_masked_lm_positions",  # : absolute position of masked tokens in input_ids
    "packed_masked_lm_ids",  # : label of masked tokens with padding as 0
    "packed_masked_lm_mask",  # : 0 if padded token. 1, 2 or 3 if masked token belongs to 1st, 2nd or 3rd sequence resp.
    "packed_next_sentence_labels",  # : 1 if next sentence, 0 otherwise
    "packed_next_sentence_mask",  # : 1 if sequence is present in pack, 0 if not
)


def expand_glob_files(files):
    result = []
    for filepath in files:
        expanded = glob.glob(filepath)
        if len(expanded) < 1:
            raise FileNotFoundError(f"Could not find file: {filepath}")
        result += expanded
    return result


class TFRecordPretrainingDataset(IterableDataset):
    """
    Preprocessed BERT pretraining dataset read from TFRecord files.

    Each datum is comprised of:
    - input_ids                  : tokens after masking
    - attention_mask             : 1 if padding token, 0 otherwise
    - token_type_ids             : sentence 0 or 1
    - masked_lm_positions        : position of masked tokens in input_ids
    - masked_lm_labels           : label of masked tokens with padding as 0
    - next_sentence_label        : 1 if next sentence, 0 otherwise

    This Dataset is compatible with multiprocessing. Each Dataloader worker
    will only read a shard of each TFRecord file, which will speed up the Dataloader
    and ensure no worker loads the same data as another worker. You are strongly
    advised to use a large number (e.g. 64) of dataloader workers because firstly,
    more workers could support high throughput, and secondly, more workers could
    give us more stochasticity and thus better convergence.


    Parameters
    ----------
    files: List of TFRecord files containing the preprocessed pretraining data
    shuffle: Shuffle the data?
    packed_data: Use packed data?
    """

    def __init__(self, input_files, shuffle=True, packed_data=False):
        self.files = expand_glob_files(input_files)
        self.shuffle = shuffle
        if packed_data:
            self.tfrecord_keys = TFRECORD_KEYS_PACKED
        else:
            self.tfrecord_keys = TFRECORD_KEYS
        self.reset()

    def reset(self):
        self.file_index = 0
        self.reader = iter([])

    def samples_per_file(self, filename):
        index_filename = filename.replace(".tfrecord", ".index")
        count = sum(1 for _ in open(index_filename))
        return count

    def __len__(self):
        if getattr(self, "_len", None) is None:
            pool = multiprocessing.Pool(min(multiprocessing.cpu_count(), len(self.files)))
            num_samples = pool.map(self.samples_per_file, self.files)
            pool.close()
            pool.join()
            self._len = sum(num_samples)
        return self._len

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is not None:
            if popdist.isPopdistEnvSet():
                self.worker_id = worker_info.id + worker_info.num_workers * popdist.getInstanceIndex()
                self.shard = (
                    worker_info.id + worker_info.num_workers * popdist.getInstanceIndex(),
                    worker_info.num_workers * popdist.getNumInstances(),
                )
            else:
                self.worker_id = worker_info.id
                self.shard = worker_info.id, worker_info.num_workers
        else:
            self.shard = None
        self.reset()
        if self.shuffle:
            np.random.shuffle(self.files)
        return self

    def __next__(self):
        try:
            datum = next(self.reader)
        except StopIteration:
            if self.file_index >= len(self.files):
                raise StopIteration
            self.reader = tfrecord_loader(
                self.files[self.file_index],
                self.files[self.file_index].replace(".tfrecord", ".index"),
                list(self.tfrecord_keys),
                self.shard,
            )
            self.file_index += 1
            datum = next(self.reader)
        datum = [datum[key] for key in self.tfrecord_keys]
        return datum


class GeneratedPretrainingDataset(Dataset):
    """
    Dataset that randomly generates mock BERT pretraining data.

    Each datum is comprised of:
    - input_ids                  : tokens after masking
    - attention_mask             : 1 if padding token, 0 otherwise
    - token_type_ids             : sentence 0 or 1
    - masked_lm_positions        : position of masked tokens in input_ids
    - masked_lm_labels           : label of masked tokens with padding as 0
    - next_sentence_label        : 1 if next sentence, 0 otherwise

    Parameters
    ----------
    vocab_size: BERT vocabulary size
    sequence_length: Sequence length
    mask_tokens: the number of mask tokens
    length: Length of generated dataset
    seed: Random seed
    packed_data: Use packed data?
    """

    def __init__(self, vocab_size, sequence_length, mask_tokens, length=1, seed=42, packed_data=False):
        self.vocab_size = vocab_size
        self.sequence_length = sequence_length
        self.mask_tokens = mask_tokens
        self.length = length
        self.seed = seed
        self.packed_data = packed_data
        self.data = self.generate_data()

    def generate_data(self):
        with torch.random.fork_rng():
            torch.manual_seed(self.seed)
            input_ids = torch.randint(0, self.vocab_size, [self.sequence_length], dtype=torch.long)
            input_mask = torch.ones_like(input_ids)
            segment_ids = torch.zeros_like(input_ids)
            masked_lm_positions = torch.randint(0, self.sequence_length, [self.mask_tokens], dtype=torch.long)
            masked_lm_ids = torch.randint(0, self.vocab_size, [self.mask_tokens], dtype=torch.long)
            if self.packed_data:
                packed_position_ids = torch.arange(self.sequence_length)
                packed_masked_lm_mask = torch.ones_like(masked_lm_ids, dtype=torch.float)
                packed_next_sentence_labels = torch.randint(0, 2, [3], dtype=torch.long)
                packed_next_sentence_mask = torch.ones_like(packed_next_sentence_labels, dtype=torch.float)
                return (
                    input_ids,
                    input_mask,
                    segment_ids,
                    packed_position_ids,
                    masked_lm_positions,
                    masked_lm_ids,
                    packed_masked_lm_mask,
                    packed_next_sentence_labels,
                    packed_next_sentence_mask,
                )
            else:
                next_sentence_labels = torch.randint(0, 2, [1], dtype=torch.long)
                return input_ids, input_mask, segment_ids, masked_lm_positions, masked_lm_ids, next_sentence_labels

    def __len__(self):
        return self.length

    def __getitem__(self, __):
        return self.data


def get_generated_datum(config):
    result = []
    dataset = GeneratedPretrainingDataset(
        config.vocab_size, config.sequence_length, config.mask_tokens, packed_data=config.packed_data
    )
    data = (dataset[i] for i in range(config.samples_per_step))
    for batches in zip(*data):
        result.append(torch.stack(batches))
    return result


class _WorkerInit:
    def __init__(self, seed):
        self.seed = seed

    def __call__(self, worker_id):
        np.random.seed((self.seed + worker_id) % np.iinfo(np.uint32).max)


def get_dataloader(config, opts):
    if config.dataset == "generated":
        dataset = GeneratedPretrainingDataset(
            config.vocab_size,
            config.sequence_length,
            config.mask_tokens,
            config.samples_per_step,
            config.random_seed,
            packed_data=config.packed_data,
        )
    elif config.dataset == "pretraining":
        dataset = TFRecordPretrainingDataset(config.input_files, packed_data=config.packed_data)
    else:
        raise RuntimeError(f"Unknown dataset '{config.dataset}', aborting.")

    loader = DataLoader(
        opts,
        dataset,
        batch_size=config.micro_batch_size,
        num_workers=config.dataloader_workers,
        worker_init_fn=_WorkerInit(config.random_seed),
        auto_distributed_partitioning=not isinstance(dataset, torch.utils.data.IterableDataset),
        mode=DataLoaderMode.AsyncRebatched if config.async_dataloader else DataLoaderMode.Sync,
    )
    return loader


if __name__ == "__main__":

    print("\nYou are executing bert_data directly.")
    print("Let's read the first input from sample dataset.")

    dataset = TFRecordPretrainingDataset(["data/sample_text.tfrecord"])
    print("dataset length: ", len(dataset), "\n")
    first = next(iter(dataset))
    named_datum = zip(
        ["input_ids", "input_mask", "segment_ids", "masked_lm_positions", "masked_lm_ids", "next_sentence_labels"],
        first,
    )
    for (name, value) in iter(named_datum):
        print(name, value.shape, value.dtype, type(value), value, "\n\n")

    print("And now, we are going to decode the tokens.\n")
    tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased", do_lower_case=True)
    print("\n\n", tokenizer.decode(first[0]), "\n\n")