def process_sample(
    sample: Dict[str, Any],
    relation_vocab: Dict[str, int],
    spacy_model: Any,
    tokenizer: Any,
) -> Tuple[Optional[Dict[str, Any]], Dict[str, int]]:
  """Processes Tacred sample and updates relation vocabulary.

  To process a raw Tacred example, we concatenate the token strings with spaces.
  The resulting text is parsed with a spacy model to find mention spans, and
  then tokenized with a BERT tokenizer. If necessary, we override some spacy
  mentions with the subj and obj Tacred mentions.

  Args:
    sample: raw Tacred sample. Needs to contain following fields: token, list of
      token strings. relation, string describing relation between subj and obj.
      subj_start, subj_end, obj_start, obj_end, starting and ending token
      indices for subj and obj (inclusive).
    relation_vocab: dictionary mapping relation strings to integer labels.
    spacy_model: spacy model used to detect mentions.
    tokenizer: BERT tokenizer.

  Returns:
    Processed Tacred sample and updated relation vocabulary.
  """

  processed_sample = {}

  relation = sample['relation']
  if relation not in relation_vocab:
    relation_vocab[relation] = len(relation_vocab)
  label = relation_vocab[relation]
  processed_sample['target'] = [label]

  words = sample['token']
  text = ' '.join(words)
  parsed_text = spacy_model(text)

  # We use spacy to parse text, identify noun chunks
  mention_char_spans = []
  # Compute subj and obj character spans
  word_lens = np.array([len(word) for word in words])
  word_char_offsets = np.cumsum(word_lens) - word_lens
  # Correct for added spaces from join
  word_char_offsets = word_char_offsets + np.arange(len(word_lens))

  def get_char_span(start_word_idx: int, end_word_idx: int) -> Tuple[int, int]:
    start_char = word_char_offsets[start_word_idx]
    # Char span is inclusive
    end_char = word_char_offsets[end_word_idx] + word_lens[end_word_idx] - 1
    assert text[start_char:end_char + 1] == ' '.join(
        words[start_word_idx:end_word_idx + 1])
    return (start_char, end_char)

  subj_char_span = get_char_span(sample['subj_start'], sample['subj_end'])
  mention_char_spans.append(subj_char_span)
  obj_char_span = get_char_span(sample['obj_start'], sample['obj_end'])
  mention_char_spans.append(obj_char_span)

  def overlaps(first_span: Tuple[int, int], second_span: Tuple[int,
                                                               int]) -> bool:

    def point_inside_span(point: int, span: Tuple[int, int]) -> bool:
      return span[0] >= point and point <= span[1]

    spans_overlap = (
        point_inside_span(first_span[0], second_span) or
        point_inside_span(first_span[1], second_span) or
        point_inside_span(second_span[0], first_span) or
        point_inside_span(second_span[1], first_span))

    return spans_overlap

  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
    char_span = (span_start_char, span_end_char)
    # Append only if does not overlap with subj or obj spans. In case spacy
    # mention annotation disagrees with tacred annotation, we want to favor
    # tacred.

    if not overlaps(char_span, subj_char_span) and not overlaps(
        char_span, obj_char_span):
      mention_char_spans.append(char_span)

  # Sort spans by start char
  start_chars = np.array([span[0] for span in mention_char_spans])
  sorted_indices = np.argsort(start_chars)
  sorted_positions = np.zeros_like(start_chars)
  sorted_positions[sorted_indices] = np.arange(len(sorted_positions))
  sorted_spans = [mention_char_spans[idx] for idx in sorted_indices]

  # Tokenize and get aligned mention positions
  _, text_ids, text_mask, mention_spans, span_indices = tokenization_utils.tokenize_with_mention_spans(
      tokenizer=tokenizer,
      sentence=text,
      spans=sorted_spans,
      max_length=FLAGS.max_length,
      add_bert_tokens=True,
      allow_truncated_spans=True,
  )

  processed_sample['text_ids'] = text_ids
  processed_sample['text_mask'] = text_mask

  # Subj and obj are the first elements of mention spans.
  subj_index = sorted_positions[0]
  obj_index = sorted_positions[1]

  # Some spans may be dropped by the BERT tokenizer. Here we map indices in the
  # original list of spans to the one returned by the tokenizer.
  reverse_span_indices = {
      original_idx: tokenized_idx
      for tokenized_idx, original_idx in enumerate(span_indices)
  }

  # Skip if subj or obj dropped.
  if (subj_index not in reverse_span_indices or
      obj_index not in reverse_span_indices):
    return None, relation_vocab

  final_subj_index = reverse_span_indices[subj_index]
  final_obj_index = reverse_span_indices[obj_index]

  processed_sample['subject_mention_indices'] = final_subj_index
  processed_sample['object_mention_indices'] = final_obj_index

  mention_spans = np.array(mention_spans)
  mention_start_positions = mention_spans[:, 0]
  mention_end_positions = mention_spans[:, 1]

  mention_start_positions = mention_start_positions[:FLAGS.max_mentions]
  mention_end_positions = mention_end_positions[:FLAGS.max_mentions]

  # We should not be truncating subject or object mentions
  assert final_subj_index < FLAGS.max_mentions
  assert final_obj_index < FLAGS.max_mentions

  mention_pad_shape = (0, FLAGS.max_mentions - len(mention_start_positions))

  mention_mask = np.ones(len(mention_start_positions), dtype=np.int64)
  mention_mask = np.pad(mention_mask, mention_pad_shape, mode='constant')
  mention_start_positions = np.pad(
      mention_start_positions, mention_pad_shape, mode='constant')
  mention_end_positions = np.pad(
      mention_end_positions, mention_pad_shape, mode='constant')

  processed_sample['mention_start_positions'] = mention_start_positions
  processed_sample['mention_end_positions'] = mention_end_positions
  processed_sample['mention_mask'] = mention_mask

  return processed_sample, relation_vocab
Processes Tacred sample and updates relation vocabulary.

  To process a raw Tacred example, we concatenate the token strings with spaces.
  The resulting text is parsed with a spacy model to find mention spans, and
  then tokenized with a BERT tokenizer. If necessary, we override some spacy
  mentions with the subj and obj Tacred mentions.

  Args:
    sample: raw Tacred sample. Needs to contain following fields: token, list of
      token strings. relation, string describing relation between subj and obj.
      subj_start, subj_end, obj_start, obj_end, starting and ending token
      indices for subj and obj (inclusive).
    relation_vocab: dictionary mapping relation strings to integer labels.
    spacy_model: spacy model used to detect mentions.
    tokenizer: BERT tokenizer.

  Returns:
    Processed Tacred sample and updated relation vocabulary.
  
language/mentionmemory/data/prepare_tacred.py