def BertSingleSentenceInputs(batch,  # pylint: disable=invalid-name
                             labeled=True,
                             cls_id=101,
                             sep_id=102):
  """Prepares inputs for BERT: add [SEP], [CLS] and create embeddings."""
  if labeled:
    for sent1, label in batch:
      value_vector = np.concatenate(([cls_id], sent1, [sep_id]))
      segment_embs = np.zeros(sent1.shape[0] + 2, dtype=np.int32)
      yield value_vector, segment_embs, segment_embs, label, np.int32(1)
  else:
    for (sent1,) in batch:  # row is a tuple with 1 element
      value_vector = np.concatenate(([cls_id], sent1, [sep_id]))
      segment_embs = np.zeros(sent1.shape[0] + 2, dtype=np.int32)
      yield value_vector, segment_embs, segment_embs
Prepares inputs for BERT: add [SEP], [CLS] and create embeddings.
trax/data/tf_inputs.py