code: def get_predictions(
      args,
      dataloader,
      biencoder_params,
      samples,
      nns,
      dists,
      mention_scores,
      cand_scores,
      pred_mention_bounds,
      id2title,
      threshold=-2.9,
      mention_threshold=-0.6931,
  ):
      """
      Arguments:
          args, dataloader, biencoder_params, samples, nns, dists, pred_mention_bounds
      Returns:
          all_entity_preds,
          num_correct_weak, num_correct_strong, num_predicted, num_gold,
          num_correct_weak_from_input_window, num_correct_strong_from_input_window, num_gold_from_input_window
      """
  
      # save biencoder predictions and print precision/recalls
      num_correct_weak = 0
      num_correct_strong = 0
      num_predicted = 0
      num_gold = 0
      num_correct_weak_from_input_window = 0
      num_correct_strong_from_input_window = 0
      num_gold_from_input_window = 0
      all_entity_preds = []
  
      f = errors_f = None
      if getattr(args, "save_preds_dir", None) is not None:
          save_biencoder_file = os.path.join(
              args.save_preds_dir, "biencoder_outs.jsonl"
          )
          f = open(save_biencoder_file, "w")
          errors_f = open(
              os.path.join(
                  args.save_preds_dir, "biencoder_errors.jsonl"
              ),
              "w",
          )
  
      # nns (List[Array[int]]) [(num_pred_mentions, cands_per_mention) x exs])
      # dists (List[Array[float]]) [(num_pred_mentions, cands_per_mention) x exs])
      # pred_mention_bounds (List[Array[int]]) [(num_pred_mentions, 2) x exs]
      # cand_scores (List[Array[float]]) [(num_pred_mentions, cands_per_mention) x exs])
      # mention_scores (List[Array[float]]) [(num_pred_mentions,) x exs])
      for batch_num, batch_data in enumerate(dataloader):
          batch_context = batch_data[0]
          if len(batch_data) > 1:
              (
                  _,
                  batch_cands,
                  batch_label_ids,
                  batch_mention_idxs,
                  batch_mention_idx_masks,
              ) = batch_data
          for b in range(len(batch_context)):
              i = batch_num * biencoder_params["eval_batch_size"] + b
              sample = samples[i]
              input_context = batch_context[b][
                  batch_context[b] != 0
              ].tolist()  # filter out padding
  
              # (num_pred_mentions, cands_per_mention)
              scores = (
                  dists[i]
                  if args.threshold_type == "joint"
                  else cand_scores[i]
              )
              cands_mask = scores[:, 0] == scores[:, 0]
              pred_entity_list = nns[i][cands_mask]
              if len(pred_entity_list) > 0:
                  e_id = pred_entity_list[0]
              distances = scores[cands_mask]
              # (num_pred_mentions, 2)
              entity_mention_bounds_idx = pred_mention_bounds[i][
                  cands_mask
              ]
              utterance = sample["text"]
  
              if args.threshold_type == "joint":
                  # THRESHOLDING
                  assert utterance is not None
                  top_mentions_mask = distances[:, 0] > threshold
              elif args.threshold_type == "top_entity_by_mention":
                  top_mentions_mask = (
                      mention_scores[i] > mention_threshold
                  )
              elif (
                  args.threshold_type == "thresholded_entity_by_mention"
              ):
                  top_mentions_mask = (distances[:, 0] > threshold) & (
                      mention_scores[i] > mention_threshold
                  )
  
              _, sort_idxs = torch.tensor(
                  distances[:, 0][top_mentions_mask]
              ).sort(descending=True)
              # cands already sorted by score
              all_pred_entities = pred_entity_list[:, 0][
                  top_mentions_mask
              ]
              e_mention_bounds = entity_mention_bounds_idx[
                  top_mentions_mask
              ]
              chosen_distances = distances[:, 0][top_mentions_mask]
              if len(all_pred_entities) >= 2:
                  all_pred_entities = all_pred_entities[sort_idxs]
                  e_mention_bounds = e_mention_bounds[sort_idxs]
                  chosen_distances = chosen_distances[sort_idxs]
  
              # prune mention overlaps
              e_mention_bounds_pruned = []
              all_pred_entities_pruned = []
              chosen_distances_pruned = []
              mention_masked_utterance = np.zeros(len(input_context))
              # ensure well-formed-ness, prune overlaps
              # greedily pick highest scoring, then prune all overlapping
              for idx, mb in enumerate(e_mention_bounds):
                  mb[
                      1
                  ] += 1  # prediction was inclusive, now make exclusive
                  # check if in existing mentions
                  if (
                      args.threshold_type != "top_entity_by_mention"
                      and mention_masked_utterance[mb[0] : mb[1]].sum()
                      >= 1
                  ):
                      continue
                  e_mention_bounds_pruned.append(mb)
                  all_pred_entities_pruned.append(
                      all_pred_entities[idx]
                  )
                  chosen_distances_pruned.append(
                      float(chosen_distances[idx])
                  )
                  mention_masked_utterance[mb[0] : mb[1]] = 1
  
              input_context = input_context[1:-1]  # remove BOS and sep
              pred_triples = [
                  (
                      str(all_pred_entities_pruned[j]),
                      int(e_mention_bounds_pruned[j][0])
                      - 1,  # -1 for BOS
                      int(e_mention_bounds_pruned[j][1]) - 1,
                  )
                  for j in range(len(all_pred_entities_pruned))
              ]
  
              entity_results = {
                  "id": sample["id"],
                  "text": sample["text"],
                  "scores": chosen_distances_pruned,
              }
  
              if "label_id" in sample:
                  # Get LABELS
                  input_mention_idxs = batch_mention_idxs[b][
                      batch_mention_idx_masks[b]
                  ].tolist()
                  input_label_ids = batch_label_ids[b][
                      batch_label_ids[b] != -1
                  ].tolist()
                  assert len(input_label_ids) == len(input_mention_idxs)
                  gold_mention_bounds = [
                      sample["text"][ment[0] - 10 : ment[0]]
                      + "["
                      + sample["text"][ment[0] : ment[1]]
                      + "]"
                      + sample["text"][ment[1] : ment[1] + 10]
                      for ment in sample["mentions"]
                  ]
  
                  # GET ALIGNED MENTION_IDXS (input is slightly different to model) between ours and gold labels -- also have to account for BOS
                  gold_input = sample["tokenized_text_ids"]
                  # return first instance of my_input in gold_input
                  for my_input_start in range(len(gold_input)):
                      if (
                          gold_input[my_input_start] == input_context[0]
                          and gold_input[
                              my_input_start : my_input_start
                              + len(input_context)
                          ]
                          == input_context
                      ):
                          break
  
                  # add alignment factor (my_input_start) to predicted mention triples
                  pred_triples = [
                      (
                          triple[0],
                          triple[1] + my_input_start,
                          triple[2] + my_input_start,
                      )
                      for triple in pred_triples
                  ]
                  gold_triples = [
                      (
                          str(sample["label_id"][j]),
                          sample["tokenized_mention_idxs"][j][0],
                          sample["tokenized_mention_idxs"][j][1],
                      )
                      for j in range(len(sample["label_id"]))
                  ]
                  (
                      num_overlap_weak,
                      num_overlap_strong,
                  ) = entity_linking_tp_with_overlap(
                      gold_triples, pred_triples
                  )
                  num_correct_weak += num_overlap_weak
                  num_correct_strong += num_overlap_strong
                  num_predicted += len(all_pred_entities_pruned)
                  num_gold += len(sample["label_id"])
  
                  # compute number correct given the input window
                  pred_input_window_triples = [
                      (
                          str(all_pred_entities_pruned[j]),
                          int(e_mention_bounds_pruned[j][0]),
                          int(e_mention_bounds_pruned[j][1]),
                      )
                      for j in range(len(all_pred_entities_pruned))
                  ]
                  gold_input_window_triples = [
                      (
                          str(input_label_ids[j]),
                          input_mention_idxs[j][0],
                          input_mention_idxs[j][1] + 1,
                      )
                      for j in range(len(input_label_ids))
                  ]
                  (
                      num_overlap_weak_window,
                      num_overlap_strong_window,
                  ) = entity_linking_tp_with_overlap(
                      gold_input_window_triples,
                      pred_input_window_triples,
                  )
                  num_correct_weak_from_input_window += (
                      num_overlap_weak_window
                  )
                  num_correct_strong_from_input_window += (
                      num_overlap_strong_window
                  )
                  num_gold_from_input_window += len(input_mention_idxs)
  
                  entity_results.update(
                      {
                          "pred_tuples_string": [
                              [
                                  id2title[triple[0]],
                                  tokenizer.decode(
                                      sample["tokenized_text_ids"][
                                          triple[1] : triple[2]
                                      ]
                                  ),
                              ]
                              for triple in pred_triples
                          ],
                          "gold_tuples_string": [
                              [
                                  id2title[triple[0]],
                                  tokenizer.decode(
                                      sample["tokenized_text_ids"][
                                          triple[1] : triple[2]
                                      ]
                                  ),
                              ]
                              for triple in gold_triples
                          ],
                          "pred_triples": pred_triples,
                          "gold_triples": gold_triples,
                          "tokens": input_context,
                      }
                  )
  
                  if errors_f is not None and (
                      num_overlap_weak != len(gold_triples)
                      or num_overlap_weak != len(pred_triples)
                  ):
                      errors_f.write(json.dumps(entity_results) + "\n")
              else:
                  entity_results.update(
                      {
                          "pred_tuples_string": [
                              [
                                  id2title[triple[0]],
                                  tokenizer.decode(
                                      input_context[
                                          triple[1] : triple[2]
                                      ]
                                  ),
                              ]
                              for triple in pred_triples
                          ],
                          "pred_triples": pred_triples,
                          "tokens": input_context,
                      }
                  )
  
              all_entity_preds.append(entity_results)
              if f is not None:
                  f.write(json.dumps(entity_results) + "\n")
  
      if f is not None:
          f.close()
          errors_f.close()
      return (
          all_entity_preds,
          num_correct_weak,
          num_correct_strong,
          num_predicted,
          num_gold,
          num_correct_weak_from_input_window,
          num_correct_strong_from_input_window,
          num_gold_from_input_window,
      )
  
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 7 for calling a global function with example usage: get_predictions(args, dataloader, biencoder_params, samples, nns, dists, mention_scores, cand_scores, pred_mention_bounds, id2title, threshold, mention_threshold) and returns: all_entity_preds,  num_correct_weak,  num_correct_strong,  num_predicted,  num_gold,  num_correct_weak_from_input_window,  num_correct_strong_from_input_window,  num_gold_from_input_window