def mask_query(self, query):
    """Applies masking to a query.

    Args:
      query: a Query object.

    Returns:
      A dict with the attributes described below.

    Raises:
      MaskingError: if the specified mask spans cannot be properly applied.

    The returned dict contains the following:
      token_ids_after_masking (list[int]): list of token IDs, including [MASK]
        tokens.
      masked_token_ids (list[int]): the original token ID of the tokens that
        have been masked out.
      mask_indices (list[int]): the token offset for all the mask tokens present
        in token_ids_after_masking.
    """
    # Tokenize.
    token_ids = [token.id for token in query.tokens]

    # Convert mask offsets to token offsets.
    mask_indices = []
    for mask_start, mask_stop in query.mask_spans:
      for token_idx, token in enumerate(query.tokens):
        # Only tokens that are strictly within a mask_span are masked.
        if token.start >= mask_start and token.stop <= mask_stop:
          mask_indices.append(token_idx)

    if query.mask_spans and not mask_indices:
      raise MaskingError('Masks do not cover any tokens.')

    # Deduplicate and sort.
    mask_indices = sorted(set(mask_indices))

    # Make sure that masked region does not get truncated in either the query
    # query representation or the joint representation.
    if mask_indices:
      if mask_indices[-1] >= self.query_seq_len - 2:
        raise MaskingError('Masks would be truncated.')

    # Mask out the specified indices.
    token_ids_after_masking = list(token_ids)  # Make a copy, then modify it.
    masked_token_ids = []
    for mask_idx in mask_indices:
      token_ids_after_masking[mask_idx] = self.tokenizer.mask_id
      masked_token_ids.append(token_ids[mask_idx])

    return {
        'token_ids_after_masking': token_ids_after_masking,
        'masked_token_ids': masked_token_ids,
        'mask_indices': mask_indices,
    }
Applies masking to a query.

    Args:
      query: a Query object.

    Returns:
      A dict with the attributes described below.

    Raises:
      MaskingError: if the specified mask spans cannot be properly applied.

    The returned dict contains the following:
      token_ids_after_masking (list[int]): list of token IDs, including [MASK]
        tokens.
      masked_token_ids (list[int]): the original token ID of the tokens that
        have been masked out.
      mask_indices (list[int]): the token offset for all the mask tokens present
        in token_ids_after_masking.
    
language/realm/featurization.py