code: def _run_biencoder(
      args,
      biencoder,
      dataloader,
      candidate_encoding,
      samples,
      num_cand_mentions=50,
      num_cand_entities=10,
      device="cpu",
      sample_to_all_context_inputs=None,
      threshold=0.0,
      indexer=None,
  ):
      """
      Returns: tuple
          labels (List[int]) [(max_num_mentions_gold) x exs]: gold labels -- returns None if no labels
          nns (List[Array[int]]) [(# of pred mentions, cands_per_mention) x exs]: predicted entity IDs in each example
          dists (List[Array[float]]) [(# of pred mentions, cands_per_mention) x exs]: scores of each entity in nns
          pred_mention_bounds (List[Array[int]]) [(# of pred mentions, 2) x exs]: predicted mention boundaries in each examples
          mention_scores (List[Array[float]]) [(# of pred mentions,) x exs]: mention score logit
          cand_scores (List[Array[float]]) [(# of pred mentions, cands_per_mention) x exs]: candidate score logit
      """
      biencoder.model.eval()
      biencoder_model = biencoder.model
      if hasattr(biencoder.model, "module"):
          biencoder_model = biencoder.model.module
  
      context_inputs = []
      nns = []
      dists = []
      mention_dists = []
      pred_mention_bounds = []
      mention_scores = []
      cand_scores = []
      sample_idx = 0
      ctxt_idx = 0
      label_ids = None
      for step, batch in enumerate(tqdm(dataloader)):
          context_input = batch[0].to(device)
          mask_ctxt = context_input != biencoder.NULL_IDX
          with torch.no_grad():
              context_outs = biencoder.encode_context(
                  context_input,
                  num_cand_mentions=num_cand_mentions,
                  topK_threshold=threshold,
              )
              embedding_ctxt = context_outs["mention_reps"]
              left_align_mask = context_outs["mention_masks"]
              chosen_mention_logits = context_outs["mention_logits"]
              chosen_mention_bounds = context_outs["mention_bounds"]
  
              """
              GET TOP CANDIDATES PER MENTION
              """
              # (all_pred_mentions_batch, embed_dim)
              embedding_ctxt = embedding_ctxt[left_align_mask]
              if indexer is None:
                  try:
                      cand_logits, _, _ = biencoder.score_candidate(
                          context_input,
                          None,
                          text_encs=embedding_ctxt,
                          cand_encs=candidate_encoding.to(device),
                      )
                      # DIM (all_pred_mentions_batch, num_cand_entities); (all_pred_mentions_batch, num_cand_entities)
                      (
                          top_cand_logits_shape,
                          top_cand_indices_shape,
                      ) = cand_logits.topk(
                          num_cand_entities, dim=-1, sorted=True
                      )
                  except:
                      # for memory savings, go through one chunk of candidates at a time
                      SPLIT_SIZE = 1000000
                      done = False
                      while not done:
                          top_cand_logits_list = []
                          top_cand_indices_list = []
                          max_chunk = int(
                              len(candidate_encoding) / SPLIT_SIZE
                          )
                          for chunk_idx in range(max_chunk):
                              try:
                                  # DIM (num_total_mentions, num_cand_entities); (num_total_mention, num_cand_entities)
                                  (
                                      top_cand_logits,
                                      top_cand_indices,
                                  ) = embedding_ctxt.mm(
                                      candidate_encoding[
                                          chunk_idx
                                          * SPLIT_SIZE : (chunk_idx + 1)
                                          * SPLIT_SIZE
                                      ]
                                      .to(device)
                                      .t()
                                      .contiguous()
                                  ).topk(
                                      10, dim=-1, sorted=True
                                  )
                                  top_cand_logits_list.append(
                                      top_cand_logits
                                  )
                                  top_cand_indices_list.append(
                                      top_cand_indices
                                      + chunk_idx * SPLIT_SIZE
                                  )
                                  if (
                                      len(
                                          (
                                              top_cand_indices_list[
                                                  chunk_idx
                                              ]
                                              < 0
                                          ).nonzero()
                                      )
                                      > 0
                                  ):
                                      import pdb
  
                                      pdb.set_trace()
                              except:
                                  SPLIT_SIZE = int(SPLIT_SIZE / 2)
                                  break
                          if len(top_cand_indices_list) == max_chunk:
                              # DIM (num_total_mentions, num_cand_entities); (num_total_mentions, num_cand_entities) -->
                              #       top_top_cand_indices_shape indexes into top_cand_indices
                              (
                                  top_cand_logits_shape,
                                  top_top_cand_indices_shape,
                              ) = torch.cat(
                                  top_cand_logits_list, dim=-1
                              ).topk(
                                  num_cand_entities, dim=-1, sorted=True
                              )
                              # make indices index into candidate_encoding
                              # DIM (num_total_mentions, max_chunk*num_cand_entities)
                              all_top_cand_indices = torch.cat(
                                  top_cand_indices_list, dim=-1
                              )
                              # DIM (num_total_mentions, num_cand_entities)
                              top_cand_indices_shape = (
                                  all_top_cand_indices.gather(
                                      -1, top_top_cand_indices_shape
                                  )
                              )
                              done = True
              else:
                  # DIM (all_pred_mentions_batch, num_cand_entities); (all_pred_mentions_batch, num_cand_entities)
                  (
                      top_cand_logits_shape,
                      top_cand_indices_shape,
                  ) = indexer.search_knn(
                      embedding_ctxt.cpu().numpy(), num_cand_entities
                  )
                  top_cand_logits_shape = torch.tensor(
                      top_cand_logits_shape
                  ).to(embedding_ctxt.device)
                  top_cand_indices_shape = torch.tensor(
                      top_cand_indices_shape
                  ).to(embedding_ctxt.device)
  
              # DIM (bs, max_num_pred_mentions, num_cand_entities)
              top_cand_logits = torch.zeros(
                  chosen_mention_logits.size(0),
                  chosen_mention_logits.size(1),
                  top_cand_logits_shape.size(-1),
              ).to(
                  top_cand_logits_shape.device,
                  top_cand_logits_shape.dtype,
              )
              top_cand_logits[left_align_mask] = top_cand_logits_shape
              top_cand_indices = torch.zeros(
                  chosen_mention_logits.size(0),
                  chosen_mention_logits.size(1),
                  top_cand_indices_shape.size(-1),
              ).to(
                  top_cand_indices_shape.device,
                  top_cand_indices_shape.dtype,
              )
              top_cand_indices[left_align_mask] = top_cand_indices_shape
  
              """
              COMPUTE FINAL SCORES FOR EACH CAND-MENTION PAIR + PRUNE USING IT
              """
              # Has NAN for impossible mentions...
              # log p(entity && mb) = log [p(entity|mention bounds) * p(mention bounds)] = log p(e|mb) + log p(mb)
              # DIM (bs, max_num_pred_mentions, num_cand_entities)
              scores = (
                  torch.log_softmax(top_cand_logits, -1)
                  + torch.sigmoid(
                      chosen_mention_logits.unsqueeze(-1)
                  ).log()
              )
  
              """
              DON'T NEED TO RESORT BY NEW SCORE -- DISTANCE PRESERVING (largest entity score still be largest entity score)
              """
  
              for idx in range(len(batch[0])):
                  # [(seqlen) x exs] <= (bsz, seqlen)
                  context_inputs.append(
                      context_input[idx][mask_ctxt[idx]]
                      .data.cpu()
                      .numpy()
                  )
                  # [(max_num_mentions, cands_per_mention) x exs] <= (bsz, max_num_mentions=num_cand_mentions, cands_per_mention)
                  nns.append(
                      top_cand_indices[idx][left_align_mask[idx]]
                      .data.cpu()
                      .numpy()
                  )
                  # [(max_num_mentions, cands_per_mention) x exs] <= (bsz, max_num_mentions=num_cand_mentions, cands_per_mention)
                  dists.append(
                      scores[idx][left_align_mask[idx]]
                      .data.cpu()
                      .numpy()
                  )
                  # [(max_num_mentions, 2) x exs] <= (bsz, max_num_mentions=num_cand_mentions, 2)
                  pred_mention_bounds.append(
                      chosen_mention_bounds[idx][left_align_mask[idx]]
                      .data.cpu()
                      .numpy()
                  )
                  # [(max_num_mentions,) x exs] <= (bsz, max_num_mentions=num_cand_mentions)
                  mention_scores.append(
                      chosen_mention_logits[idx][left_align_mask[idx]]
                      .data.cpu()
                      .numpy()
                  )
                  # [(max_num_mentions, cands_per_mention) x exs] <= (bsz, max_num_mentions=num_cand_mentions, cands_per_mention)
                  cand_scores.append(
                      top_cand_logits[idx][left_align_mask[idx]]
                      .data.cpu()
                      .numpy()
                  )
  
      return (
          nns,
          dists,
          pred_mention_bounds,
          mention_scores,
          cand_scores,
      )
  
imports: import argparse
  import json
  import sys
  from elq.index.faiss_indexer import DenseFlatIndexer, DenseHNSWFlatIndexer, DenseIVFFlatIndexer
  import logging
  import torch
  import numpy
  from colorama import init
  from termcolor import colored
  import torch.nn.functional
  import blink.ner
  from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
  from elq.biencoder.biencoder import BiEncoderRanker, load_biencoder, to_bert_input
  from elq.biencoder.data_process import process_mention_data, get_context_representation_single_mention, get_candidate_representation
  import elq.candidate_ranking.utils
  import math
  from elq.vcg_utils.measures import entity_linking_tp_with_overlap
  from elq.biencoder.utils import batch_reshape_mask_left
  import os
  import sys
  from tqdm import tqdm
  import pdb
  import time
  from transformers import BertTokenizer
label_desc: Use image node_id 6 for calling a global function with example usage: _run_biencoder(args, biencoder, dataloader, candidate_encoding, samples, num_cand_mentions, num_cand_entities, device, sample_to_all_context_inputs, threshold, indexer) and returns: nns,  dists,  pred_mention_bounds,  mention_scores,  cand_scores