Delete pretrain_text_dataloader.py

pretrain_text_dataloader.py

@@ -1,226 +0,0 @@
-# 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)