def add_entity_tokens(
    text_ids: tf.Tensor,
    text_mask: tf.Tensor,
    mention_mask: tf.Tensor,
    mention_batch_positions: tf.Tensor,
    mention_start_positions: tf.Tensor,
    mention_end_positions: tf.Tensor,
    new_length: int,
    mlm_target_positions: Optional[tf.Tensor] = None,
    mlm_target_weights: Optional[tf.Tensor] = None,
    entity_start_token_id: int = default_values.ENTITY_START_TOKEN,
    entity_end_token_id: int = default_values.ENTITY_END_TOKEN,
) -> Dict[str, tf.Tensor]:
  """Adds entity start / end tokens around mentions.

  Inserts `entity_start_token_id` and `entity_end_token_id` tokens around each
  mention and update mention_start_positions / mention_end_positions to point
  to these tokens.

  New text length will be `new_length` and texts will be truncated if nessesary.
  If a mention no longer fits into the new text, its mask (`mention_mask`) will
  be set to 0.

  The function can also update MLM position and weights (`mlm_target_positions`
  and `mlm_target_weights`) if these arguments are provided. Similarly to
  mentions, if MLM position no longer fits into the new text, its mask
  (`mlm_target_weights`) will be set to 0.

  Args:
    text_ids: [seq_length] tensor with token ids.
    text_mask: [seq_length] tensor with 1s for tokens and 0 for padding.
    mention_mask: [n_mentions] mask indicating whether a mention is a padding.
    mention_batch_positions: [n_mentions] sample ID of a mention in the batch.
    mention_start_positions: [n_mentions] position of a mention first token
      within a sample.
    mention_end_positions: [n_mentions] position of a mention last token within
      a sample.
    new_length: new length of text after entity tokens will be added.
    mlm_target_positions: [batch_size, max_mlm_targets] positions of tokens to
      be used for MLM task.
    mlm_target_weights: [batch_size, max_mlm_targets] mask indicating whether
      `mlm_target_positions` is a padding.
    entity_start_token_id: token to be used as entity start token.
    entity_end_token_id: token to be used as entity end token.

  Returns:
    New text_ids and text_mask, updated mentions positions including
    mention_start_positions, mention_end_positions and mention_mask.
    Returns updated mlm_target_positions and mlm_target_weights if they were
    provided as arguments.
  """
  batch_size = tf.shape(text_ids)[0]
  old_length = tf.shape(text_ids)[1]
  new_shape = (batch_size, new_length)

  mentions_fit_mask = compute_which_mentions_fit_with_entity_tokens(
      mention_mask,
      mention_batch_positions,
      mention_start_positions,
      mention_end_positions,
      batch_size,
      old_length,
      new_length,
  )
  # Ignore mentions that does not fit into new texts.
  new_mention_mask = mention_mask * mentions_fit_mask
  mention_start_positions = mention_start_positions * new_mention_mask
  mention_end_positions = mention_end_positions * new_mention_mask

  positions = compute_positions_shift_with_entity_tokens(
      new_mention_mask, mention_batch_positions, mention_start_positions,
      mention_end_positions, batch_size, old_length)

  def get_2d_index(positions: tf.Tensor) -> tf.Tensor:
    return _get_2d_index(mention_batch_positions, positions)

  def get_new_positions(old_positions: tf.Tensor) -> tf.Tensor:
    index_2d = get_2d_index(old_positions)
    return tf.gather_nd(positions, index_2d)

  new_mention_start_positions = get_new_positions(mention_start_positions) - 1
  new_mention_start_positions = new_mention_start_positions * new_mention_mask
  new_mention_end_positions = get_new_positions(mention_end_positions) + 1
  new_mention_end_positions = new_mention_end_positions * new_mention_mask

  if mlm_target_positions is not None:
    if mlm_target_weights is None:
      raise ValueError('`mlm_target_weights` must be specified if '
                       '`mlm_target_positions` is provided.')
    mlm_target_positions = tf.gather(
        positions, mlm_target_positions, batch_dims=1)
    mlm_target_positions_within_len = tf.less(mlm_target_positions, new_length)
    mlm_target_positions_within_len = tf.cast(mlm_target_positions_within_len,
                                              mlm_target_weights.dtype)
    mlm_target_weights = mlm_target_weights * mlm_target_positions_within_len
    # Zero-out positions for pad MLM targets
    mlm_target_positions = mlm_target_positions * mlm_target_weights

  # Cut texts that are longer than `new_length`
  text_within_new_length = tf.less(positions, new_length)
  text_ids = text_ids * tf.cast(text_within_new_length, text_ids.dtype)
  text_mask = text_mask * tf.cast(text_within_new_length, text_mask.dtype)
  positions = tf.minimum(positions, new_length - 1)

  # Prepare 2D index for tokens positions in the next text_ids and text_mask.
  # Note that we use flat 2D index and flat values
  # (e.g. `tf.reshape(text_ids, [-1])`) since `tf.scatter_nd` does not support
  # batch dimension.
  batch_positions = _batched_range(old_length, batch_size, 1, positions.dtype)
  batch_positions = tf.reshape(batch_positions, [-1])
  text_index_2d = _get_2d_index(batch_positions, tf.reshape(positions, [-1]))

  new_text_ids = tf.scatter_nd(text_index_2d, tf.reshape(text_ids, [-1]),
                               new_shape)
  new_text_mask = tf.scatter_nd(text_index_2d, tf.reshape(text_mask, [-1]),
                                new_shape)

  # Insert entity start / end tokens into the new text_ids and text_mask.
  new_mention_start_positions_2d = get_2d_index(new_mention_start_positions)
  new_mention_end_positions_2d = get_2d_index(new_mention_end_positions)

  new_text_ids = tf.tensor_scatter_nd_add(
      new_text_ids, new_mention_start_positions_2d,
      new_mention_mask * entity_start_token_id)
  new_text_ids = tf.tensor_scatter_nd_add(
      new_text_ids, new_mention_end_positions_2d,
      new_mention_mask * entity_end_token_id)

  new_mention_mask = tf.cast(new_mention_mask, dtype=text_mask.dtype)
  new_text_mask = tf.tensor_scatter_nd_add(new_text_mask,
                                           new_mention_start_positions_2d,
                                           new_mention_mask)
  new_text_mask = tf.tensor_scatter_nd_add(new_text_mask,
                                           new_mention_end_positions_2d,
                                           new_mention_mask)

  features = {
      'text_ids': new_text_ids,
      'text_mask': new_text_mask,
      'mention_start_positions': new_mention_start_positions,
      'mention_end_positions': new_mention_end_positions,
      'mention_mask': new_mention_mask,
  }

  if mlm_target_positions is not None:
    features['mlm_target_weights'] = mlm_target_weights
    features['mlm_target_positions'] = mlm_target_positions

  return features
Adds entity start / end tokens around mentions.

  Inserts `entity_start_token_id` and `entity_end_token_id` tokens around each
  mention and update mention_start_positions / mention_end_positions to point
  to these tokens.

  New text length will be `new_length` and texts will be truncated if nessesary.
  If a mention no longer fits into the new text, its mask (`mention_mask`) will
  be set to 0.

  The function can also update MLM position and weights (`mlm_target_positions`
  and `mlm_target_weights`) if these arguments are provided. Similarly to
  mentions, if MLM position no longer fits into the new text, its mask
  (`mlm_target_weights`) will be set to 0.

  Args:
    text_ids: [seq_length] tensor with token ids.
    text_mask: [seq_length] tensor with 1s for tokens and 0 for padding.
    mention_mask: [n_mentions] mask indicating whether a mention is a padding.
    mention_batch_positions: [n_mentions] sample ID of a mention in the batch.
    mention_start_positions: [n_mentions] position of a mention first token
      within a sample.
    mention_end_positions: [n_mentions] position of a mention last token within
      a sample.
    new_length: new length of text after entity tokens will be added.
    mlm_target_positions: [batch_size, max_mlm_targets] positions of tokens to
      be used for MLM task.
    mlm_target_weights: [batch_size, max_mlm_targets] mask indicating whether
      `mlm_target_positions` is a padding.
    entity_start_token_id: token to be used as entity start token.
    entity_end_token_id: token to be used as entity end token.

  Returns:
    New text_ids and text_mask, updated mentions positions including
    mention_start_positions, mention_end_positions and mention_mask.
    Returns updated mlm_target_positions and mlm_target_weights if they were
    provided as arguments.
  
language/mentionmemory/utils/mention_preprocess_utils.py