def featurize_document_tf(self, title_token_ids, body_token_ids):
    """Featurizes a document using only TF ops (for non-eager execution).

    Args:
      title_token_ids: a 1-D int Tensor.
      body_token_ids: a 1-D int Tensor.

    Returns:
      A dict with the following structure:
        input_ids: [candidate_seq_len] int32 Tensor
        input_mask: [candidate_seq_len] int32 Tensor
        segment_ids: [candidate_seq_len] int32 Tensor
    """
    # Total number of tokens representing the doc, including special tokens.
    sequence_len = self.candidate_seq_len

    # Possibly truncate title, reserving at least 3 spots for special tokens.
    title_token_ids = title_token_ids[:sequence_len - 3]
    title_len = tf.size(title_token_ids)

    # Calculate how many spots left for body text.
    max_body_len = tf.maximum(sequence_len - 3 - title_len, 0)

    # Possibly truncate the body text.
    body_token_ids = body_token_ids[:max_body_len]
    body_len = tf.size(body_token_ids)

    # Zero-padding to occupy whatever space remains.
    pair_len = title_len + body_len
    padding = tf.zeros([sequence_len - 3 - pair_len], tf.int32)

    # Put title, body and padding together.
    input_ids = tf.concat([
        [self.tokenizer.cls_id],
        title_token_ids,
        [self.tokenizer.sep_id],
        body_token_ids,
        [self.tokenizer.sep_id],
        padding], 0)

    # input_mask indicates non-pad tokens.
    input_mask = tf.concat([tf.ones([pair_len + 3], tf.int32), padding], 0)

    # segment_ids optionally distinguish title from body.
    if self.separate_candidate_segments:
      segment_ids = tf.concat([
          tf.zeros([title_len + 2], tf.int32),
          tf.ones([body_len + 1], tf.int32), padding
      ], 0)
    else:
      segment_ids = tf.concat([tf.zeros([pair_len + 3], tf.int32), padding], 0)

    # Add static shape annotations.
    input_ids.set_shape([sequence_len])
    input_mask.set_shape([sequence_len])
    segment_ids.set_shape([sequence_len])

    return {'input_ids': input_ids,
            'input_mask': input_mask,
            'segment_ids': segment_ids}
Featurizes a document using only TF ops (for non-eager execution).

    Args:
      title_token_ids: a 1-D int Tensor.
      body_token_ids: a 1-D int Tensor.

    Returns:
      A dict with the following structure:
        input_ids: [candidate_seq_len] int32 Tensor
        input_mask: [candidate_seq_len] int32 Tensor
        segment_ids: [candidate_seq_len] int32 Tensor
    
language/realm/featurization.py