pretrain_text_dataloader.py

# Copyright 2024 The TensorFlow Authors. 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.

"""Loads text dataset for the BERT pretraining task."""
import dataclasses
from typing import List, Mapping, Optional, Text

import tensorflow as tf, tf_keras
import tensorflow_text as tf_text

from official.common import dataset_fn
from official.core import config_definitions as cfg
from official.core import input_reader
from official.nlp.data import data_loader
from official.nlp.data import data_loader_factory
from official.nlp.modeling.ops import segment_extractor


@dataclasses.dataclass
class BertPretrainTextDataConfig(cfg.DataConfig):
  """Data config for BERT pretraining task (tasks/masked_lm) from text."""
  input_path: str = ""
  doc_batch_size: int = 8
  global_batch_size: int = 512
  is_training: bool = True
  seq_length: int = 512
  max_predictions_per_seq: int = 76
  use_next_sentence_label: bool = True
  # The name of the text feature fields. The text features will be
  # concatenated in order.
  # Note: More than 1 field name is not compatible with NSP.
  text_field_names: Optional[List[str]] = dataclasses.field(
      default_factory=lambda: ["text"])
  vocab_file_path: str = ""
  masking_rate: float = 0.15
  use_whole_word_masking: bool = False
  file_type: str = "tfrecord"


_CLS_TOKEN = b"[CLS]"
_SEP_TOKEN = b"[SEP]"
_MASK_TOKEN = b"[MASK]"
_NUM_OOV_BUCKETS = 1
# Accounts for [CLS] and 2 x [SEP] tokens
_NUM_SPECIAL_TOKENS = 3


@data_loader_factory.register_data_loader_cls(BertPretrainTextDataConfig)
class BertPretrainTextDataLoader(data_loader.DataLoader):
  """A class to load text dataset for BERT pretraining task."""

  def __init__(self, params):
    """Inits `BertPretrainTextDataLoader` class.

    Args:
      params: A `BertPretrainTextDataConfig` object.
    """
    if len(params.text_field_names) > 1 and params.use_next_sentence_label:
      raise ValueError("Currently there is no support for more than text field "
                       "while generating next sentence labels.")

    self._params = params
    self._seq_length = params.seq_length
    self._max_predictions_per_seq = params.max_predictions_per_seq
    self._use_next_sentence_label = params.use_next_sentence_label
    self._masking_rate = params.masking_rate
    self._use_whole_word_masking = params.use_whole_word_masking

    lookup_table_init = tf.lookup.TextFileInitializer(
        params.vocab_file_path,
        key_dtype=tf.string,
        key_index=tf.lookup.TextFileIndex.WHOLE_LINE,
        value_dtype=tf.int64,
        value_index=tf.lookup.TextFileIndex.LINE_NUMBER)
    self._vocab_lookup_table = tf.lookup.StaticVocabularyTable(
        lookup_table_init,
        num_oov_buckets=_NUM_OOV_BUCKETS,
        lookup_key_dtype=tf.string)

    self._cls_token = self._vocab_lookup_table.lookup(tf.constant(_CLS_TOKEN))
    self._sep_token = self._vocab_lookup_table.lookup(tf.constant(_SEP_TOKEN))
    self._mask_token = self._vocab_lookup_table.lookup(tf.constant(_MASK_TOKEN))

    # -_NUM_OOV_BUCKETS to offset unused OOV bucket.
    self._vocab_size = self._vocab_lookup_table.size() - _NUM_OOV_BUCKETS

  def _decode(self, record: tf.Tensor) -> Mapping[Text, tf.Tensor]:
    """Decodes a serialized tf.Example."""
    name_to_features = {}
    for text_field_name in self._params.text_field_names:
      name_to_features[text_field_name] = tf.io.FixedLenFeature([], tf.string)
    return tf.io.parse_single_example(record, name_to_features)

  def _tokenize(self, segments):
    """Tokenize the input segments."""
    # Tokenize segments
    tokenizer = tf_text.BertTokenizer(
        self._vocab_lookup_table, token_out_type=tf.int64)

    if self._use_whole_word_masking:
      # tokenize the segments which should have the shape:
      # [num_sentence, (num_words), (num_wordpieces)]
      segments = [tokenizer.tokenize(s) for s in segments]
    else:
      # tokenize the segments and merge out the token dimension so that each
      # segment has the shape: [num_sentence, (num_wordpieces)]
      segments = [tokenizer.tokenize(s).merge_dims(-2, -1) for s in segments]

    # Truncate inputs
    trimmer = tf_text.WaterfallTrimmer(
        self._seq_length - _NUM_SPECIAL_TOKENS, axis=-1)
    truncated_segments = trimmer.trim(segments)

    # Combine segments, get segment ids and add special tokens
    return tf_text.combine_segments(
        truncated_segments,
        start_of_sequence_id=self._cls_token,
        end_of_segment_id=self._sep_token)

  def _bert_preprocess(self, record: Mapping[str, tf.Tensor]):
    """Parses raw tensors into a dict of tensors to be consumed by the model."""
    if self._use_next_sentence_label:
      input_text = record[self._params.text_field_names[0]]
      # Split sentences
      sentence_breaker = tf_text.RegexSplitter()
      sentences = sentence_breaker.split(input_text)

      # Extract next-sentence-prediction labels and segments
      next_or_random_segment, is_next = (
          segment_extractor.get_next_sentence_labels(sentences))
      # merge dims to change shape from [num_docs, (num_segments)] to
      # [total_num_segments]
      is_next = is_next.merge_dims(-2, -1)

      # construct segments with shape [(num_sentence)]
      segments = [
          sentences.merge_dims(-2, -1),
          next_or_random_segment.merge_dims(-2, -1)
      ]
    else:
      segments = [record[name] for name in self._params.text_field_names]

    segments_combined, segment_ids = self._tokenize(segments)

    # Dynamic masking
    item_selector = tf_text.RandomItemSelector(
        self._max_predictions_per_seq,
        selection_rate=self._masking_rate,
        unselectable_ids=[self._cls_token, self._sep_token],
        shuffle_fn=(tf.identity if self._params.deterministic else None))
    values_chooser = tf_text.MaskValuesChooser(
        vocab_size=self._vocab_size, mask_token=self._mask_token)
    masked_input_ids, masked_lm_positions, masked_lm_ids = (
        tf_text.mask_language_model(
            segments_combined,
            item_selector=item_selector,
            mask_values_chooser=values_chooser,
        ))

    # Pad out to fixed shape and get input mask.
    seq_lengths = {
        "input_word_ids": self._seq_length,
        "input_type_ids": self._seq_length,
        "masked_lm_positions": self._max_predictions_per_seq,
        "masked_lm_ids": self._max_predictions_per_seq,
    }
    model_inputs = {
        "input_word_ids": masked_input_ids,
        "input_type_ids": segment_ids,
        "masked_lm_positions": masked_lm_positions,
        "masked_lm_ids": masked_lm_ids,
    }
    padded_inputs_and_mask = tf.nest.map_structure(tf_text.pad_model_inputs,
                                                   model_inputs, seq_lengths)
    model_inputs = {
        k: padded_inputs_and_mask[k][0] for k in padded_inputs_and_mask
    }
    model_inputs["masked_lm_weights"] = tf.cast(
        padded_inputs_and_mask["masked_lm_ids"][1], tf.float32)
    model_inputs["input_mask"] = padded_inputs_and_mask["input_word_ids"][1]

    if self._use_next_sentence_label:
      model_inputs["next_sentence_labels"] = is_next

    for name in model_inputs:
      t = model_inputs[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      model_inputs[name] = t

    return model_inputs

  def load(self, input_context: Optional[tf.distribute.InputContext] = None):
    """Returns a tf.dataset.Dataset."""

    def _batch_docs(dataset, input_context):
      per_core_doc_batch_size = (
          input_context.get_per_replica_batch_size(self._params.doc_batch_size)
          if input_context else self._params.doc_batch_size)
      return dataset.batch(per_core_doc_batch_size)

    reader = input_reader.InputReader(
        params=self._params,
        dataset_fn=dataset_fn.pick_dataset_fn(self._params.file_type),
        decoder_fn=self._decode if self._params.input_path else None,
        transform_and_batch_fn=_batch_docs
        if self._use_next_sentence_label else None,
        postprocess_fn=self._bert_preprocess)
    transformed_inputs = reader.read(input_context)
    per_core_example_batch_size = (
        input_context.get_per_replica_batch_size(self._params.global_batch_size)
        if input_context else self._params.global_batch_size)
    batched_inputs = transformed_inputs.unbatch().batch(
        per_core_example_batch_size, self._params.drop_remainder)
    return batched_inputs.prefetch(tf.data.experimental.AUTOTUNE)