Function bert_format

language/realm/featurization.py:555–611

Combines multiple sequences of tokens into a single sequence for BERT. The format is: [CLS] first seq [SEP] second seq [SEP] third_seq [SEP] ... [SEP] If the combined tokens exceed max_seq_len: - We try to keep as much of the first sequence as possible. - If room remains, we try to keep

(input_seqs, max_seq_len, cls_id, sep_id)

def bert_format(input_seqs, max_seq_len, cls_id, sep_id):
  """Combines multiple sequences of tokens into a single sequence for BERT.

  The format is:
  [CLS] first seq [SEP] second seq [SEP] third_seq [SEP] ... [SEP]

  If the combined tokens exceed max_seq_len:
  - We try to keep as much of the first sequence as possible.
  - If room remains, we try to keep as much of the second sequence.
  - ... and so on.
  - We always make room for all of the [CLS] and [SEP] tokens.

  Args:
    input_seqs (list): a list of sequences, where each seq is a list of ints.
    max_seq_len (int): max sequence length.
    cls_id (int): integer ID for the [CLS] token.
    sep_id (int): integer ID for the [SEP] token.

  Returns:
    an instance of BertInputs, where each attribute is an int32
    Tensor with shape [max_seq_len].
  """
  # Overall token budget, after reserving space for special tokens.
  token_budget = max_seq_len - 1 - len(input_seqs)
  if token_budget < 0:
    raise ValueError(
        'max_seq_len not large enough to include all special tokens.')

  input_ids = [cls_id]
  segment_ids = [0]
  input_mask = [1]

  for segment_id, raw_input_seq in enumerate(input_seqs):
    # Truncate to stay within the remaining token budget.
    input_seq = raw_input_seq[:token_budget]
    input_seq_len = len(input_seq)

    input_ids.extend(input_seq)
    input_ids.append(sep_id)

    segment_ids.extend([segment_id] * (input_seq_len + 1))
    input_mask.extend([1] * (input_seq_len + 1))

    # Subtract from budget.
    token_budget -= input_seq_len

  assert len(input_ids) == len(input_mask)
  assert len(input_ids) == len(segment_ids)

  # Pad with zeroes up to max_seq_len, and convert to TF Tensor.
  as_tensor = lambda arr: tf.constant(truncate_or_pad(arr, max_seq_len, 0))
  input_ids = as_tensor(input_ids)
  input_mask = as_tensor(input_mask)
  segment_ids = as_tensor(segment_ids)

  return BertInputs(
      token_ids=input_ids, mask=input_mask, segment_ids=segment_ids)