code: def get_context_representation_multiple_mentions_idxs(
      sample,
      tokenizer,
      max_seq_length,
      mention_key,
      context_key,
      ent_start_token,
      ent_end_token,
  ):
      """
      Also cuts out mentions beyond that context window
  
      ASSUMES MENTION_IDXS ARE SORTED!!!!
  
      Returns:
          List of mention bounds that are [inclusive, exclusive) (make both inclusive later)
          NOTE: 2nd index of mention bound may be outside of max_seq_length-range (must deal with later)
      """
      mention_idxs = sample["tokenized_mention_idxs"]
      input_ids = sample["tokenized_text_ids"]
  
      # sort mentions / entities / everything associated
      # [[orig_index, [start, end]], ....] --> sort by start, then end
      sort_tuples = [
          [i[0], i[1]]
          for i in sorted(
              enumerate(mention_idxs), key=lambda x: (x[1][0], x[1][1])
          )
      ]
      if [tup[1] for tup in sort_tuples] != mention_idxs:
          orig_idx_to_sort_idx = {
              itm[0]: i for i, itm in enumerate(sort_tuples)
          }
          assert [tup[1] for tup in sort_tuples] == sort_mentions(
              mention_idxs, orig_idx_to_sort_idx
          )
          mention_idxs = [tup[1] for tup in sort_tuples]
          sample["tokenized_mention_idxs"] = mention_idxs
          do_sort(sample, orig_idx_to_sort_idx)
          # TODO SORT EVERYTHING
  
      # fit leftmost mention, then all of the others that can reasonably fit...
      all_mention_spans_range = [
          mention_idxs[0][0],
          mention_idxs[-1][1],
      ]
      while (
          all_mention_spans_range[1] - all_mention_spans_range[0] + 2
          > max_seq_length
      ):
          if len(mention_idxs) == 1:
              # don't cut further
              assert (
                  mention_idxs[0][1] - mention_idxs[0][0] + 2
                  > max_seq_length
              )
              # truncate mention
              mention_idxs[0][1] = (
                  max_seq_length + mention_idxs[0][0] - 2
              )
          else:
              # cut last mention
              mention_idxs = mention_idxs[: len(mention_idxs) - 1]
          all_mention_spans_range = [
              mention_idxs[0][0],
              mention_idxs[-1][1],
          ]
  
      context_left = input_ids[: all_mention_spans_range[0]]
      all_mention_tokens = input_ids[
          all_mention_spans_range[0] : all_mention_spans_range[1]
      ]
      context_right = input_ids[all_mention_spans_range[1] :]
  
      left_quota = (max_seq_length - len(all_mention_tokens)) // 2 - 1
      right_quota = (
          max_seq_length - len(all_mention_tokens) - left_quota - 2
      )
      left_add = len(context_left)
      right_add = len(context_right)
      if (
          left_add <= left_quota
      ):  # tokens left to add <= quota ON THE LEFT
          if (
              right_add > right_quota
          ):  # add remaining quota to right quota
              right_quota += left_quota - left_add
      else:
          if (
              right_add <= right_quota
          ):  # tokens left to add <= quota ON THE RIGHT
              left_quota += (
                  right_quota - right_add
              )  # add remaining quota to left quota
  
      if left_quota <= 0:
          left_quota = -len(
              context_left
          )  # cut entire list (context_left = [])
      if right_quota <= 0:
          right_quota = 0  # cut entire list (context_right = [])
      input_ids_window = (
          context_left[-left_quota:]
          + all_mention_tokens
          + context_right[:right_quota]
      )
  
      # shift mention_idxs
      if len(input_ids) <= max_seq_length - 2:
          try:
              assert input_ids == input_ids_window
          except:
              import pdb
  
              pdb.set_trace()
      else:
          assert input_ids != input_ids_window
          cut_from_left = len(context_left) - len(
              context_left[-left_quota:]
          )
          if cut_from_left > 0:
              # must shift mention_idxs
              for c in range(len(mention_idxs)):
                  mention_idxs[c] = [
                      mention_idxs[c][0] - cut_from_left,
                      mention_idxs[c][1] - cut_from_left,
                  ]
  
      input_ids_window = [101] + input_ids_window + [102]
      tokens = tokenizer.convert_ids_to_tokens(input_ids_window)
  
      # +1 for CLS token
      mention_idxs = [
          [mention[0] + 1, mention[1] + 1] for mention in mention_idxs
      ]
  
      # input_ids = tokenizer.convert_tokens_to_ids(input_ids_window)
      padding = [0] * (max_seq_length - len(input_ids_window))
      input_ids_window += padding
      assert len(input_ids_window) == max_seq_length
  
      return {
          "tokens": tokens,
          "ids": input_ids_window,
          "mention_idxs": mention_idxs,
          # "pruned_ents": [1 for i in range(len(all_mentions)) if i < len(mention_idxs) else 0],  # pruned last N entities, TODO change if changed
      }
  
file_path: BLINK/elq/biencoder/data_process.py
imports: import os
  import json
  import logging
  import torch
  from tqdm import tqdm, trange
  from torch.utils.data import DataLoader, TensorDataset
  from pytorch_transformers.tokenization_bert import BertTokenizer
  from blink.biencoder.zeshel_utils import world_to_id
  from elq.common.params import ENT_START_TAG, ENT_END_TAG, ENT_TITLE_TAG
label_desc: Use image node_id 7 for calling a global function with example usage: get_context_representation_multiple_mentions_idxs(sample, tokenizer, max_seq_length, mention_key, context_key, ent_start_token, ent_end_token) and returns: dict