def gen_mention_pretraining_sample(
    text_length: int,
    n_mentions: int,
    n_linked_mentions: int,
    max_length: int = 100,
    vocab_size: int = 100,
    entity_vocab_size: int = 1000,
    mention_size: int = 2,
) -> Dict[str, np.ndarray]:
  """Generate test raw decoded input for mention pre-training pipeline."""

  text_pad_shape = (0, max_length - text_length)
  text_ids = np.random.randint(
      low=1, high=vocab_size, size=text_length, dtype=np.int64)
  text_ids = np.pad(text_ids, pad_width=text_pad_shape, mode='constant')
  text_mask = np.pad(
      np.ones(shape=text_length, dtype=np.int64),
      pad_width=text_pad_shape,
      mode='constant')

  mention_start_positions = np.random.choice(
      text_length // mention_size, size=n_mentions,
      replace=False) * mention_size
  mention_start_positions.sort()
  mention_end_positions = mention_start_positions + mention_size - 1

  dense_span_starts = np.zeros(shape=max_length, dtype=np.int64)
  dense_span_starts[mention_start_positions] = 1

  dense_span_ends = np.zeros(shape=max_length, dtype=np.int64)
  dense_span_ends[mention_end_positions] = 1

  linked_mention_indices = np.arange(n_linked_mentions)

  linked_mention_position_slices = [
      np.arange(mention_start_positions[idx], mention_end_positions[idx] + 1)
      for idx in linked_mention_indices
  ]

  if n_linked_mentions > 0:
    dense_linked_mention_positions = np.concatenate(
        linked_mention_position_slices)
  else:
    dense_linked_mention_positions = np.arange(0)

  linked_mention_ids = np.random.randint(
      low=1, high=entity_vocab_size, size=len(linked_mention_indices))

  dense_mention_mask = np.zeros(shape=max_length, dtype=np.int64)
  dense_mention_mask[dense_linked_mention_positions] = 1
  dense_mention_ids = np.zeros(shape=max_length, dtype=np.int64)
  for idx, position_slice in enumerate(linked_mention_position_slices):
    dense_mention_ids[position_slice] = linked_mention_ids[idx]

  dense_answer_mask = np.ones_like(dense_mention_mask)

  raw_example = {
      'text_ids': text_ids,
      'text_mask': text_mask,
      'dense_span_starts': dense_span_starts,
      'dense_span_ends': dense_span_ends,
      'dense_mention_mask': dense_mention_mask,
      'dense_mention_ids': dense_mention_ids,
      'dense_answer_mask': dense_answer_mask,
  }
  return raw_example
Generate test raw decoded input for mention pre-training pipeline.
language/mentionmemory/utils/test_utils.py