def _gen_input(self, text_len: int, n_mentions: int):
    """Generate synthetic inputs to test masking."""
    text_ids = np.pad(
        np.random.randint(1, self.vocab_size, size=text_len),
        (0, self.max_len - text_len),
        mode='constant')
    text_ids = text_ids.astype(np.int64)
    text_mask = np.pad(
        np.ones(text_len, dtype=np.int64), (0, self.max_len - text_len),
        mode='constant')
    mention_start_positions = np.random.choice(
        max((text_len - self.mention_size) // self.mention_size, 0),
        size=n_mentions,
        replace=False) * self.mention_size
    mention_start_positions.sort()
    mention_end_positions = mention_start_positions + self.mention_size - 1
    mention_end_positions = np.minimum(mention_end_positions, self.max_len - 1)
    dense_span_starts = np.zeros_like(text_ids)
    dense_span_starts[mention_start_positions] = 1
    dense_span_ends = np.zeros_like(text_ids)
    dense_span_ends[mention_end_positions] = 1
    dense_mention_mask = np.zeros(self.max_len, dtype=np.int64)
    dense_mention_ids = np.zeros(self.max_len, dtype=np.int64)
    for idx in range(n_mentions):
      mention_slice = slice(mention_start_positions[idx],
                            mention_end_positions[idx] + 1)
      dense_mention_mask[mention_slice] = 1
      dense_mention_ids[mention_slice] = np.random.randint(1, 1000)

    # This feature is only for `test_process_batchwise_mention_targets`
    dense_is_masked = np.random.randint(2, size=self.max_len)
    return {
        'text_ids': text_ids,
        'text_mask': text_mask,
        'dense_span_starts': dense_span_starts,
        'dense_span_ends': dense_span_ends,
        'mention_start_positions': mention_start_positions,
        'mention_end_positions': mention_end_positions,
        'dense_mention_ids': dense_mention_ids,
        'dense_mention_mask': dense_mention_mask,
        'dense_is_masked': dense_is_masked,
    }
Generate synthetic inputs to test masking.
language/mentionmemory/utils/mention_preprocess_utils_test.py