Function bert_preprocess

language/conpono/reconstruct/preprocess.py:477–575

Pre-processes a text tuple into BERT format. Args: text_segments (Tensor): 1-D Tensor of un-tokenized text segments (string). These will be concatenated into a single sequence. vocab_table (StaticHashTable): a map from string to integer. max_seq_length (int): max # of wordpiece

(text_segments,
                    vocab_table,
                    max_seq_length,
                    max_predictions_per_seq,
                    tokenize_fn=None,
                    mask_rate=0.15,
                    do_lower_case=True,
                    cls_token=_CLS_TOKEN,
                    sep_token=_SEP_TOKEN,
                    mask_token=_MASK_TOKEN)

def bert_preprocess(text_segments,
                    vocab_table,
                    max_seq_length,
                    max_predictions_per_seq,
                    tokenize_fn=None,
                    mask_rate=0.15,
                    do_lower_case=True,
                    cls_token=_CLS_TOKEN,
                    sep_token=_SEP_TOKEN,
                    mask_token=_MASK_TOKEN):
  """Pre-processes a text tuple into BERT format.

  Args:
    text_segments (Tensor): 1-D Tensor of un-tokenized text segments (string).
      These will be concatenated into a single sequence.
    vocab_table (StaticHashTable): a map from string to integer.
    max_seq_length (int): max # of wordpiece tokens, including special tokens.
    max_predictions_per_seq (int): max # of masked positions for entire seq.
    tokenize_fn: a function that performs a tokenization on a 1D string Tensor
      of untokenized strings. By default this uses wordpiece tokenization.
    mask_rate (float): percentage of tokens to mask out. If 0, no masking is
      performed.
    do_lower_case (bool): whether to lowercase text or not. Default is True.
    cls_token (unicode): token representing CLS
    sep_token (unicode): token representing SEP (separator)
    mask_token (unicode): token representing MASK

  Returns:
    BertInputs
  """

  if not tokenize_fn:
    # pylint: disable=g-long-lambda
    tokenize_fn = lambda text_input: wordpiece_tokenize(
        text_input=text_segments,
        vocab_table=vocab_table,
        token_out_type=tf.string,
        use_unknown_token=True,
        lower_case=do_lower_case)

  segment_tokens_grouped = tokenize_fn(text_segments)
  # TODO(kguu): check that BERT handles UNK in the same way.

  # This is a RaggedTensor of shape [batch_size, (num_wordpieces)]
  # One row for each segment. Each row is a list of wordpiece tokens.
  segment_tokens = collapse_dims(segment_tokens_grouped)

  # Truncate.
  num_special_tokens = segment_tokens.nrows() + 1
  segment_tokens_truncated = truncate_segment_tokens(
      segment_tokens, max_seq_length - num_special_tokens)

  # Add special tokens.
  segment_tokens_with_special_tokens = add_special_tokens(
      segment_tokens_truncated, cls_token, sep_token)

  # Compute segment IDs.
  segment_ids_2d = create_segment_ids(segment_tokens_with_special_tokens)