def process_sample(
    sample: Dict[str, Any],
    spacy_model: Any,
    tokenizer: Any,
) -> Optional[Dict[str, Any]]:
  """Processes CWQ sample.

  To process a raw CWQ example, we first strip special tokens and convert tokens
  back to text. Then a spacy model extracts noun chunks, and the standard BERT
  tokenizer tokenizes again. Finally we append a mask token to the text, and add
  the answer entity id as a mention at the position of the mask token.

  Args:
    sample: raw CWQ sample. Needs to contain following fields: text_ids,
      target_mention.
    spacy_model: spacy model used to extract mentions.
    tokenizer: BERT tokenizer.

  Returns:
    Processed CWQ sample.
  """
  numpy_sample = {key: value.numpy() for key, value in sample.items()}
  processed_sample = {}

  text_ids = numpy_sample['text_ids']
  target_mention = numpy_sample['target_mention']
  # Skip samples for which the answer is not present in entity vocabulary
  if target_mention == 0:
    return None

  text_ids = text_ids[text_ids.nonzero()]

  # Remove special tokens
  special_tokens_to_remove = [101, 102, 103, 104, 105, 106]
  keep_positions = np.logical_not(np.isin(text_ids, special_tokens_to_remove))
  text_ids = text_ids[keep_positions]

  def detokenize(text_ids, tokenizer):
    text_ids = text_ids[np.nonzero(text_ids)]
    text_tokens = tokenizer.convert_ids_to_tokens(text_ids)
    text_tokens = [
        t[2:] if t.startswith('##') else ' ' + t for t in text_tokens
    ]
    return ''.join(text_tokens)

  # Convert back to text
  text = detokenize(text_ids, tokenizer)

  # Annotate noun chunks
  parsed_text = spacy_model(text)
  mention_char_spans = []
  for chunk in parsed_text.noun_chunks:
    span_start_char = parsed_text[chunk.start].idx
    span_last_token = parsed_text[chunk.end - 1]
    span_end_char = span_last_token.idx + len(span_last_token.text) - 1
    mention_char_spans.append((span_start_char, span_end_char))

  # Tokenize again. Set max length to 1 less than max length as we still have to
  # add the mask token afterwards.
  _, text_ids, text_mask, mention_spans, _ = tokenization_utils.tokenize_with_mention_spans(
      tokenizer=tokenizer,
      sentence=text,
      spans=mention_char_spans,
      max_length=FLAGS.max_length - 1,
      add_bert_tokens=True,
  )
  text_ids = np.array(text_ids)
  text_ids = text_ids[np.nonzero(text_mask)]
  answer_position = len(text_ids)

  # Add mask token
  text_ids = np.append(text_ids, default_values.MASK_TOKEN)

  text_pad_shape = (0, FLAGS.max_length - len(text_ids))
  text_ids = np.pad(text_ids, text_pad_shape)
  text_mask = (text_ids > 0) * 1

  # Truncate mentions, keeping room for final mention.
  mention_spans = mention_spans[:FLAGS.max_mentions - 1]
  mention_spans.append((answer_position, answer_position))
  mention_spans = np.array(mention_spans)

  mention_start_positions = mention_spans[:, 0]
  mention_end_positions = mention_spans[:, 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_ids = np.zeros_like(text_ids)
  dense_mention_ids[answer_position] = target_mention
  dense_mention_mask = (dense_mention_ids > 0) * 1

  processed_sample['text_ids'] = text_ids
  processed_sample['text_mask'] = text_mask
  processed_sample['dense_mention_ids'] = dense_mention_ids
  processed_sample['dense_mention_mask'] = dense_mention_mask
  processed_sample['dense_span_starts'] = dense_span_starts
  processed_sample['dense_span_ends'] = dense_span_ends

  return processed_sample
Processes CWQ sample.

  To process a raw CWQ example, we first strip special tokens and convert tokens
  back to text. Then a spacy model extracts noun chunks, and the standard BERT
  tokenizer tokenizes again. Finally we append a mask token to the text, and add
  the answer entity id as a mention at the position of the mask token.

  Args:
    sample: raw CWQ sample. Needs to contain following fields: text_ids,
      target_mention.
    spacy_model: spacy model used to extract mentions.
    tokenizer: BERT tokenizer.

  Returns:
    Processed CWQ sample.
  
language/mentionmemory/data/prepare_complexwq.py