def BertDoubleSentenceInputs(batch,  # pylint: disable=invalid-name
                             labeled=True,
                             cls_id=101,
                             sep_id=102):
  """Prepares inputs for BERT models by adding [SEP] and [CLS] tokens and creating segment embeddings."""
  if labeled:
    for sent1, sent2, label in batch:
      value_vector = np.concatenate(
          ([cls_id], sent1, [sep_id], sent2, [sep_id]))

      segment_embs = np.zeros(
          sent1.shape[0] + sent2.shape[0] + 3, dtype=np.int32)
      second_sent_start = sent1.shape[0] + 2
      segment_embs[second_sent_start:] = 1
      yield value_vector, segment_embs, segment_embs, label, np.int32(1)
  else:
    for sent1, sent2 in batch:
      value_vector = np.concatenate(
          ([cls_id], sent1, [sep_id], sent2, [sep_id]))

      segment_embs = np.zeros(
          sent1.shape[0] + sent2.shape[0] + 3, dtype=np.int32)
      second_sent_start = sent1.shape[0] + 2
      segment_embs[second_sent_start:] = 1
      yield value_vector, segment_embs, segment_embs
Prepares inputs for BERT models by adding [SEP] and [CLS] tokens and creating segment embeddings.
trax/data/tf_inputs.py