code: def process_mention_data(
      samples,
      tokenizer,
      max_context_length,
      max_cand_length,
      silent,
      mention_key="mention",
      context_key="context",
      label_key="label",
      title_key="label_title",
      ent_start_token=ENT_START_TAG,
      ent_end_token=ENT_END_TAG,
      title_token=ENT_TITLE_TAG,
      debug=False,
      logger=None,
      add_mention_bounds=True,
      saved_context_dir=None,
      candidate_token_ids=None,
      params=None,
  ):
      """
      Returns /inclusive/ bounds
      """
      extra_ret_values = {}
      if saved_context_dir is not None and os.path.exists(
          os.path.join(saved_context_dir, "tensor_tuple.pt")
      ):
          data = torch.load(os.path.join(saved_context_dir, "data.pt"))
          tensor_data_tuple = torch.load(
              os.path.join(saved_context_dir, "tensor_tuple.pt")
          )
          return data, tensor_data_tuple, extra_ret_values
  
      if candidate_token_ids is None and not debug:
          candidate_token_ids = torch.load(
              params["cand_token_ids_path"]
          )
          if logger:
              logger.info("Loaded saved entities info")
          extra_ret_values["candidate_token_ids"] = candidate_token_ids
  
      processed_samples = []
  
      if debug:
          samples = samples[:200]
      if silent:
          iter_ = samples
      else:
          iter_ = tqdm(samples)
  
      use_world = True
  
      ent_start_id = tokenizer.convert_tokens_to_ids(ent_start_token)
      ent_end_id = tokenizer.convert_tokens_to_ids(ent_end_token)
      cls_token_id = tokenizer.convert_tokens_to_ids("[CLS]")
      sep_token_id = tokenizer.convert_tokens_to_ids("[SEP]")
      for idx, sample in enumerate(iter_):
          assert (
              not add_mention_bounds
          ), "Adding mention bounds, but we have multiple entities per example"
          if context_key + "_left" in sample:
              context_tokens = get_context_representation_multiple_mentions_left_right(
                  sample,
                  tokenizer,
                  max_context_length,
                  mention_key,
                  context_key,
                  ent_start_token,
                  ent_end_token,
              )
          else:
              context_tokens = (
                  get_context_representation_multiple_mentions_idxs(
                      sample,
                      tokenizer,
                      max_context_length,
                      mention_key,
                      context_key,
                      ent_start_token,
                      ent_end_token,
                  )
              )
  
          for i in range(len(context_tokens["mention_idxs"])):
              context_tokens["mention_idxs"][i][
                  1
              ] -= 1  # make bounds inclusive
  
          label = sample[label_key]
          title = sample.get(title_key)
          label_ids = sample.get("label_id")
  
          if label is None:
              label = [None]
              label_ids = [label_ids]
          # remove those that got pruned off
          if len(label) > len(context_tokens["mention_idxs"]):
              label = label[: len(context_tokens["mention_idxs"])]
              label_ids = sample["label_id"][
                  : len(context_tokens["mention_idxs"])
              ]
  
          if candidate_token_ids is not None:
              token_ids = [
                  [candidate_token_ids[label_id].tolist()]
                  for label_id in label_ids
              ]
              label_tokens = {
                  "tokens": "",
                  "ids": token_ids,
              }
          elif not params["freeze_cand_enc"]:
              label_tokens = [
                  get_candidate_representation(
                      l,
                      tokenizer,
                      max_cand_length,
                      title[i],
                  )
                  for i, l in enumerate(label)
              ]
              label_tokens = {
                  k: [
                      label_tokens[l][k]
                      for l in range(len(label_tokens))
                  ]
                  for k in label_tokens[0]
              }
          else:
              label_tokens = None
  
          if isinstance(sample["label_id"], list):
              # multiple candidates
              if len(sample["label_id"]) > len(
                  context_tokens["mention_idxs"]
              ):
                  sample["label_id"] = sample["label_id"][
                      : len(context_tokens["mention_idxs"])
                  ]
              label_idx = [int(id) for id in sample["label_id"]]
          else:
              assert isinstance(sample["label_id"], int) or isinstance(
                  sample["label_id"], str
              )
              label_idx = int(sample["label_id"])
  
          record = {
              "context": context_tokens,
          }
          if not params["freeze_cand_enc"]:
              record["label"] = label_tokens
          record["label_idx"] = label_idx
  
          if "world" in sample:
              src = sample["world"]
              src = world_to_id[src]
              record["src"] = [src]
              use_world = True
          else:
              use_world = False
  
          processed_samples.append(record)
  
      if debug and logger:
          logger.info("====Processed samples: ====")
          for sample in processed_samples[:5]:
              logger.info(
                  "Context tokens : "
                  + " ".join(sample["context"]["tokens"])
              )
              logger.info(
                  "Context ids : "
                  + " ".join([str(v) for v in sample["context"]["ids"]])
              )
              if not params["freeze_cand_encs"]:
                  logger.info(
                      "Label tokens : "
                      + " ".join(sample["label"]["tokens"])
                  )
                  logger.info(
                      "Label ids : "
                      + " ".join(
                          [str(v) for v in sample["label"]["ids"]]
                      )
                  )
              logger.info("Label_id : %d" % sample["label_idx"])
              if use_world:
                  logger.info("Src : %d" % sample["src"][0])
  
      context_vecs = torch.tensor(
          select_field(processed_samples, "context", "ids"),
          dtype=torch.long,
      )
      if logger:
          logger.info("Created context IDs vector")
      if isinstance(
          processed_samples[0]["context"]["mention_idxs"][0], int
      ):
          mention_idx_vecs = torch.tensor(
              select_field(
                  processed_samples, "context", "mention_idxs"
              ),
              dtype=torch.long,
          ).unsqueeze(1)
          mention_idx_mask = torch.ones(
              mention_idx_vecs.size(0), dtype=torch.bool
          ).unsqueeze(-1)
          if logger:
              logger.info("Created mention positions vector")
  
          if not params["freeze_cand_enc"]:
              cand_vecs = torch.tensor(
                  select_field(processed_samples, "label", "ids"),
                  dtype=torch.long,
              )
              if logger:
                  logger.info("Created candidate IDs vector")
          label_idx = torch.tensor(
              select_field(processed_samples, "label_idx"),
              dtype=torch.long,
          ).unsqueeze(-1)
          if logger:
              logger.info("Created label IDXs vector")
      else:
          (
              mention_idx_vecs,
              mention_idx_mask,
          ) = select_field_with_padding(
              processed_samples,
              "context",
              "mention_idxs",
              pad_idx=[0, 1],  # ensure is a well-formed span
          )
          # (bs, max_num_spans, 2)
          mention_idx_vecs = torch.tensor(
              mention_idx_vecs, dtype=torch.long
          )
          # (bs, max_num_spans)
          mention_idx_mask = torch.tensor(
              mention_idx_mask, dtype=torch.bool
          )
  
          if not params["freeze_cand_enc"]:
              cand_vecs, cand_mask = select_field_with_padding(
                  processed_samples,
                  "label",
                  "ids",
                  pad_idx=[[0 for _ in range(max_cand_length)]],
              )
              # (bs, max_num_spans, 1, max_cand_length)
              cand_vecs = torch.tensor(cand_vecs, dtype=torch.long)
              cand_mask = torch.tensor(cand_mask, dtype=torch.bool)
              assert (
                  cand_mask == mention_idx_mask
              ).all() or cand_mask.all()
              if logger:
                  logger.info("Created candidate IDs vector")
          else:
              cand_vecs = torch.Tensor(context_vecs.size())
  
          label_idx_vecs, label_idx_mask = select_field_with_padding(
              processed_samples, "label_idx", pad_idx=-1
          )
          # (bs, max_num_spans)
          label_idx = torch.tensor(label_idx_vecs, dtype=torch.long)
          label_idx_mask = torch.tensor(
              label_idx_mask, dtype=torch.bool
          )
          assert (
              label_idx_mask == mention_idx_mask
          ).all() or label_idx_mask.all()
          if logger:
              logger.info("Created label IDXs vector")
      # mention_idx_vecs: (bs, max_num_spans, 2), mention_idx_mask: (bs, max_num_spans)
      assert len(mention_idx_vecs.size()) == 3
      # prune mention_idx_vecs to max_context_length
      mention_idx_vecs[mention_idx_vecs >= max_context_length] = (
          max_context_length - 1
      )
  
      if use_world:
          src_vecs = torch.tensor(
              select_field(processed_samples, "src"),
              dtype=torch.long,
          )
          if logger:
              logger.info("Created source vector")
  
      data = {
          "context_vecs": context_vecs,
          "mention_idx_vecs": mention_idx_vecs,
          "cand_vecs": cand_vecs,
          "label_idx": label_idx,
      }
  
      if use_world:
          data["src"] = src_vecs
          tensor_data_tuple = (
              context_vecs,
              cand_vecs,
              src_vecs,
              label_idx,
              mention_idx_vecs,
              mention_idx_mask,
          )
      else:
          tensor_data_tuple = (
              context_vecs,
              cand_vecs,
              label_idx,
              mention_idx_vecs,
              mention_idx_mask,
          )
      # save data
      if saved_context_dir is not None and not os.path.exists(
          os.path.join(saved_context_dir, "tensor_tuple.pt")
      ):
          os.makedirs(saved_context_dir, exist_ok=True)
          torch.save(data, os.path.join(saved_context_dir, "data.pt"))
          torch.save(
              tensor_data_tuple,
              os.path.join(saved_context_dir, "tensor_tuple.pt"),
          )
      return data, tensor_data_tuple, extra_ret_values
  
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 9 for calling a global function with example usage: process_mention_data(samples, tokenizer, max_context_length, max_cand_length, silent, mention_key, context_key, label_key, title_key, ent_start_token, ent_end_token, title_token, debug, logger, add_mention_bounds, saved_context_dir, candidate_token_ids, params) and returns: data,  tensor_data_tuple,  extra_ret_values, data,  tensor_data_tuple,  extra_ret_values