def BertNextSentencePredictionInputs(dataset_name,  # pylint: disable=invalid-name
                                     data_dir=None,
                                     text_key='text',
                                     train=True,
                                     shuffle_size=50000):
  """Defines a stream for the next sentence prediction task."""
  stream = TFDS(
      dataset_name,
      data_dir=data_dir,
      tfds_preprocess_fn=functools.partial(
          t5_data().preprocessors.next_sentence_prediction,
          text_key=text_key,
          label_sentences=True,
          buffer_size=shuffle_size),
      keys=['inputs', 'targets'],
      train=train)

  def split_stream(generator=None):
    # split string with 'sentence1:' and 'sentence2:' into two separate strings
    for text, target in stream(generator):
      text_str = str(text)[:-1]  # removes last '"' which is always at the end
      sentences = text_str.split('sentence1: ')[1].split(' sentence2: ')
      if len(sentences) != 2:
        # 'sentence2:' appeared in the text and got mixed up with the label
        continue
      sent1, sent2 = sentences
      yield sent1, sent2, target == 'next'

  return split_stream
Defines a stream for the next sentence prediction task.
trax/data/tf_inputs.py