data.py

import functools
import itertools
import os
import typing

import datasets
import torch
import tokenizers
import transformers

import utils


def wt_detokenizer(string):
  string = string.replace("s '", "s'")
  string = string.replace(" @-@ ", "-")
  string = string.replace(" @,@ ", ",")
  string = string.replace(" @.@ ", ".")
  string = string.replace(" : ", ": ")
  string = string.replace(" ; ", "; ")
  string = string.replace(" . ", ". ")
  string = string.replace(" ! ", "! ")
  string = string.replace(" ? ", "? ")
  string = string.replace(" , ", ", ")
  string = string.replace(" \n", "\n")
  string = string.replace("\n ", "\n")
  string = string.replace(" 's", "'s")
  return string


def ptb_detokenizer(x):
  x = x.replace(" 's", "'s")
  x = x.replace("s ' ", "s' ")
  x = x.replace(" n't", "n't")
  x = x.replace(" \n ", "\n")
  x = x.replace("\\/", "/")
  for _ in range(10):
    x = x.replace(" N ", " 1 ")
  x = x.replace("$ 1", "$1")
  x = x.replace("# 1", "#1")
  x = x.replace("<unk>", "?")
  return x


def lm1b_detokenizer(x):
  x = x.replace('http : / / ', 'http://')
  x = x.replace('https : / / ', 'https://')
  x = x.replace(' ? ', '? ')
  x = x.replace(' ! ', '! ')
  x = x.replace(' , ', ', ')
  x = x.replace(' : ', ': ')
  x = x.replace(' ; ', '; ')
  x = x.replace(' / ', '/')
  return x


class Text8Tokenizer(transformers.PreTrainedTokenizer):
  def __init__(
    self,
    bos_token='[BOS]',
    eos_token='[EOS]',
    sep_token='[SEP]',
    cls_token='[CLS]',
    pad_token='[PAD]',
    mask_token='[MASK]',
    unk_token='[UNK]',
    **kwargs):
    self.characters = list('abcdefghijklmnopqrstuvwxyz ')
    self._vocab_str_to_int = {
      '[CLS]': 0,
      '[SEP]': 1,
      '[BOS]': 2,
      '[EOS]': 3,
      '[MASK]': 4,
      '[PAD]': 5,
      '[RESERVED]': 6,
      '[UNK]': 7,
      ** {ch: i + 8 for i, ch in enumerate(self.characters)}}
    self._vocab_int_to_str = {
      v: k for k, v in self._vocab_str_to_int.items()}
    super().__init__(
      bos_token=bos_token,
      eos_token=eos_token,
      sep_token=sep_token,
      cls_token=cls_token,
      pad_token=pad_token,
      mask_token=mask_token,
      unk_token=unk_token,
      **kwargs)

  @property
  def vocab_size(self) -> int:
    return len(self._vocab_str_to_int)

  def _tokenize(self, text: str, **kwargs) -> typing.List[str]:
    return list(text.lower())

  def _convert_token_to_id(self, token: str) -> int:
    return self._vocab_str_to_int.get(
      token, self._vocab_str_to_int['[UNK]'])

  def _convert_id_to_token(self, index: int) -> str:
    return self._vocab_int_to_str[index]

  def convert_tokens_to_string(self, tokens):
    return ''.join(tokens)

  def get_vocab(self) -> typing.Dict[str, int]:
    return self._vocab_str_to_int


def _group_texts(examples, block_size, bos, eos):
  concatenated_examples = list(itertools.chain(* examples['input_ids']))
  total_length = len(concatenated_examples)
  new_block_size = block_size - 2
  total_length = (total_length // new_block_size) * new_block_size
  result = {}
  _values = []
  _attn_masks = []
  for i in range(0, total_length, new_block_size):
    _values.append(
      [bos]
      + concatenated_examples[i : i + new_block_size]
      + [eos])
    _attn_masks.append(torch.ones(block_size))
  result['input_ids'] = _values
  result['attention_mask'] = _attn_masks
  return result


def get_tokenizer(tokenizer_name: str,
                  add_mask_token: bool = True,
                  add_reserved: bool = True):
  if tokenizer_name == 'text8':
    tokenizer = Text8Tokenizer()
  elif tokenizer_name == 'bert-base-uncased':
    tokenizer = transformers.BertTokenizer.from_pretrained(
      'bert-base-uncased')
  else:
    tokenizer = transformers.AutoTokenizer.from_pretrained(
      tokenizer_name)

  if tokenizer.bos_token is None:
    tokenizer.bos_token = tokenizer.cls_token or tokenizer.eos_token
  if tokenizer.eos_token is None:
    tokenizer.eos_token = tokenizer.sep_token or tokenizer.bos_token
  if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
  if add_mask_token and tokenizer.mask_token is None:
    tokenizer.add_special_tokens({'mask_token': '[MASK]'})

  if add_reserved:
    tokenizer.add_special_tokens({'additional_special_tokens': ['[RE]']})

  if (isinstance(tokenizer, transformers.GPT2TokenizerFast)
      or isinstance(tokenizer, transformers.GPT2Tokenizer)):
    tokenizer._tokenizer.post_processor = tokenizers.processors.BertProcessing(
      (tokenizer.bos_token, tokenizer.bos_token_id),
      (tokenizer.eos_token, tokenizer.eos_token_id))

  return tokenizer


def get_dataset(dataset_name,
                tokenizer,
                mode='train',
                cache_dir='./data_cache',
                block_size=256,
                wrap=False,
                streaming=False,
                num_proc=len(os.sched_getaffinity(0))):
  filename = f'{dataset_name}_{mode}_bs{block_size}_{"wrap" if wrap else "nowrap"}.dat'
  _path = os.path.join(cache_dir, filename)
  os.makedirs(cache_dir, exist_ok=True)

  if utils.fsspec_exists(_path):
    return datasets.load_from_disk(_path).with_format('torch')

  if dataset_name == 'wikitext103':
    dataset = datasets.load_dataset(
      'wikitext', name='wikitext-103-raw-v1',
      cache_dir=cache_dir)
  elif dataset_name == 'wikitext2':
    dataset = datasets.load_dataset(
      'wikitext', name='wikitext-2-raw-v1',
      cache_dir=cache_dir)
  elif dataset_name == 'ptb':
    dataset = datasets.load_dataset(
      'ptb_text_only', cache_dir=cache_dir)
  elif dataset_name == 'openwebtext-train':
    dataset = datasets.load_dataset(
      'openwebtext', split='train[:-100000]',
      cache_dir=cache_dir, streaming=streaming)
  elif dataset_name == 'openwebtext-valid':
    dataset = datasets.load_dataset(
      'openwebtext', split='train[-100000:]',
      cache_dir=cache_dir, streaming=streaming)
  else:
    dataset = datasets.load_dataset(
      dataset_name,
      cache_dir=cache_dir,
      streaming=streaming)

  if isinstance(dataset, datasets.IterableDatasetDict):
    data = dataset[mode]
  elif isinstance(dataset, datasets.IterableDataset):
    data = dataset
  elif isinstance(dataset, datasets.Dataset):
    data = dataset
  else:
    data = dataset[mode]

  if dataset_name.startswith('wikitext'):
    detokenizer = wt_detokenizer
  elif dataset_name == 'ptb':
    detokenizer = ptb_detokenizer
  elif dataset_name == 'lm1b':
    detokenizer = lm1b_detokenizer
  else:
    detokenizer = None

  def _apply_detokenizer(detokenizer):
    def detok(text):
      for i, t in enumerate(text, 0):
        text[i] = detokenizer(t)
      return text
    return detok

  EOS = tokenizer.encode(tokenizer.eos_token)[0]
  BOS = tokenizer.encode(tokenizer.bos_token)[0]

  def preprocess_and_tokenize(example):
    text = example['text'] if 'text' in example else example.get('article', '')
    if detokenizer is not None:
      text = _apply_detokenizer(detokenizer)(text)
    tokenizer.padding_side = 'right'
    tokenizer.truncation_side = 'right'

    if wrap:
      tokens = tokenizer(text,
                         add_special_tokens=False,
                         return_attention_mask=False,
                         return_token_type_ids=False)
      tokens = {'input_ids':
                [t + [EOS] for t in tokens['input_ids']]}
    else:
      tokens = tokenizer(text,
                         max_length=block_size,
                         padding='max_length',
                         truncation=True,
                         add_special_tokens=True,
                         return_attention_mask=True,
                         return_token_type_ids=False)
    return tokens

  if streaming:
    tokenized_dataset = data.map(
      preprocess_and_tokenize,
      batched=True,
      desc='Tokenizing')
  else:
    tokenized_dataset = data.map(
      preprocess_and_tokenize,
      batched=True,
      num_proc=num_proc,
      load_from_cache_file=True,
      desc='Tokenizing')

  columns_to_remove = [c for c in tokenized_dataset.column_names
                       if c not in ['input_ids', 'attention_mask']]
  tokenized_dataset = tokenized_dataset.remove_columns(columns_to_remove)

  if not wrap:
    tokenized_dataset.save_to_disk(_path)
    return tokenized_dataset.with_format('torch')

  group_texts = functools.partial(
    _group_texts, block_size=block_size, bos=BOS, eos=EOS)
  if streaming:
    chunked_dataset = tokenized_dataset.map(
      group_texts,
      batched=True,
      desc='Grouping')
  else:
    chunked_dataset = tokenized_dataset.map(
      group_texts,
      batched=True,
      num_proc=num_proc,
      load_from_cache_file=True,
      desc='Grouping')
    chunked_dataset.save_to_disk(_path)
  chunked_dataset = chunked_dataset.with_format('torch')
  return chunked_dataset


class MaskedDataset(torch.utils.data.Dataset):
  def __init__(self, dataset, masker):
    self.dataset = dataset
    self.masker = masker

  def __len__(self):
    try:
      return len(self.dataset)
    except TypeError:
      return 0

  def __getitem__(self, idx):
    sample = self.dataset[idx]
    masked = self.masker(
      sample['input_ids'],
      sample.get('attention_mask', None))
    return masked


def build_dataloaders(tokenizer,
                      masker,
                      dataset_name,
                      cache_dir,
                      max_length,
                      batch_size,
                      eval_batch_size,
                      num_workers=4,
                      wrap=False,
                      streaming=False):
  if streaming:
    raise NotImplementedError(
      'Streaming dataloaders are not supported in this factorization setup.')
  train_name = dataset_name
  valid_name = dataset_name
  train_mode = 'train'
  if dataset_name.startswith('openwebtext'):
    train_name = 'openwebtext-train'
    valid_name = 'openwebtext-valid'
    train_mode = 'train'

  train_set = get_dataset(
    train_name,
    tokenizer,
    mode=train_mode,
    cache_dir=cache_dir,
    block_size=max_length,
    wrap=wrap,
    streaming=streaming)
  if dataset_name == 'text8':
    validation_split = 'test'
  elif dataset_name in ['openwebtext-train', 'openwebtext-valid']:
    validation_split = 'validation' if dataset_name != 'openwebtext-valid' else 'train'
  else:
    validation_split = 'validation'
  valid_set = get_dataset(
    valid_name,
    tokenizer,
    mode=validation_split,
    cache_dir=cache_dir,
    block_size=max_length,
    wrap=wrap,
    streaming=streaming)

  train_loader = torch.utils.data.DataLoader(
    MaskedDataset(train_set, masker),
    batch_size=batch_size,
    shuffle=not streaming,
    num_workers=num_workers,
    pin_memory=True)
  valid_loader = torch.utils.data.DataLoader(
    MaskedDataset(valid_set, masker),
    batch_size=eval_batch_size,
    shuffle=False,
    num_workers=num_workers,
    pin_memory=True)
  return train_loader, valid_loader