docstr: if gold_mention_bounds specified, selects according to gold_mention_bounds,
  otherwise selects according to top-scoring mentions
  
  Returns: Dictionary
      mention_reps: torch.FloatTensor (bsz, max_num_pred_mentions, embed_dim): mention embeddings
      mention_masks: torch.BoolTensor (bsz, max_num_pred_mentions): mention padding mask
      mention_bounds: torch.LongTensor (bsz, max_num_pred_mentions, 2)
      (
      mention_logits: torch.FloatTensor (bsz, max_num_pred_mentions): mention scores/logits
      all_mention_mask: torch.BoolTensor ((bsz, all_cand_mentions)
      all_mention_logits: torch.FloatTensor (bsz, all_cand_mentions): all mention scores/logits
      all_mention_bounds: torch.LongTensor (bsz, all_cand_mentions, 2): all mention bounds
      )
docstr_tok: ["if","gold_mention_bounds","specified",",","selects","according","to","gold_mention_bounds",",","otherwise","selects","according","to","top-scoring","mentions","Returns",":","Dictionary","mention_reps",":","torch.FloatTensor","(","bsz",",","max_num_pred_mentions",",","embed_dim",")",":","mention","embeddings","mention_masks",":","torch.BoolTensor","(","bsz",",","max_num_pred_mentions",")",":","mention","padding","mask","mention_bounds",":","torch.LongTensor","(","bsz",",","max_num_pred_mentions",",","2",")","(","mention_logits",":","torch.FloatTensor","(","bsz",",","max_num_pred_mentions",")",":","mention","scores\/logits","all_mention_mask",":","torch.BoolTensor","(","(","bsz",",","all_cand_mentions",")","all_mention_logits",":","torch.FloatTensor","(","bsz",",","all_cand_mentions",")",":","all","mention","scores\/logits","all_mention_bounds",":","torch.LongTensor","(","bsz",",","all_cand_mentions",",","2",")",":","all","mention","bounds",")"]
code: def encode_context(
      self,
      cands,
      gold_mention_bounds=None,
      gold_mention_bounds_mask=None,
      num_cand_mentions=50,
      topK_threshold=-4.5,
      get_mention_scores=True,
  ):
      """
      if gold_mention_bounds specified, selects according to gold_mention_bounds,
      otherwise selects according to top-scoring mentions
  
      Returns: Dictionary
          mention_reps: torch.FloatTensor (bsz, max_num_pred_mentions, embed_dim): mention embeddings
          mention_masks: torch.BoolTensor (bsz, max_num_pred_mentions): mention padding mask
          mention_bounds: torch.LongTensor (bsz, max_num_pred_mentions, 2)
          (
          mention_logits: torch.FloatTensor (bsz, max_num_pred_mentions): mention scores/logits
          all_mention_mask: torch.BoolTensor ((bsz, all_cand_mentions)
          all_mention_logits: torch.FloatTensor (bsz, all_cand_mentions): all mention scores/logits
          all_mention_bounds: torch.LongTensor (bsz, all_cand_mentions, 2): all mention bounds
          )
      """
      token_idx_cands, segment_idx_cands, mask_cands = to_bert_input(
          cands, self.NULL_IDX
      )
      context_outs, _ = self.model(
          token_idx_cands,
          segment_idx_cands,
          mask_cands,
          None,
          None,
          None,
          gold_mention_bounds=gold_mention_bounds,
          gold_mention_bounds_mask=gold_mention_bounds_mask,
          num_cand_mentions=num_cand_mentions,
          topK_threshold=topK_threshold,
          get_mention_scores=get_mention_scores,
      )
      if context_outs["mention_dims"].size(0) <= 1:
          for key in context_outs:
              if "all" in key or key == "mention_dims":
                  continue
              context_outs[key] = context_outs[key].view(
                  [context_outs["mention_dims"][0, 0], -1]
                  + list(context_outs[key].size()[1:])
              )
          return context_outs
  
      """
          Reshape to (bs, num_mentions, *), iterating across GPUs
          """
  
      def init_tensor(shape, dtype, init_value):
          return init_value * torch.ones(shape).to(
              dtype=dtype, device=context_outs["mention_dims"].device
          )
  
      bs = cands.size(0)
      n_pred_mentions = context_outs["mention_dims"][:, 1].max()
      context_outs_reshape = {}
      for key in context_outs:
          if "all" in key or key == "mention_dims":
              context_outs_reshape[key] = context_outs[key]
              continue
          # (bsz, max_num_pred_mentions, *)
          context_outs_reshape[key] = init_tensor(
              [bs, n_pred_mentions]
              + list(context_outs[key].size()[1:]),
              context_outs[key].dtype,
              -float("inf") if "logit" in key else 0,
          )
  
      for idx in range(len(context_outs["mention_dims"])):
          # reshape
          gpu_bs = context_outs["mention_dims"][idx, 0]
          b_width = context_outs["mention_dims"][idx, 1]
  
          start_idx = (
              context_outs["mention_dims"][:idx, 0]
              * context_outs["mention_dims"][:idx, 1]
          ).sum()
          end_idx = start_idx + b_width * gpu_bs
  
          s_reshape = context_outs["mention_dims"][:idx, 0].sum()
          e_reshape = s_reshape + gpu_bs
          for key in context_outs_reshape:
              if "all" in key or key == "mention_dims":
                  continue
              if len(context_outs[key].size()) == 1:
                  target_tensor = context_outs[key][
                      start_idx:end_idx
                  ].view(gpu_bs, b_width)
              else:
                  target_tensor = context_outs[key][
                      start_idx:end_idx
                  ].view(gpu_bs, b_width, -1)
              context_outs_reshape[key][
                  s_reshape:e_reshape, :b_width
              ] = target_tensor
  
      return context_outs_reshape
  
imports: import os
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from tqdm import tqdm
  from collections import OrderedDict
  from pytorch_transformers.modeling_bert import BertPreTrainedModel, BertConfig, BertModel
  from pytorch_transformers.tokenization_bert import BertTokenizer
  from elq.common.ranker_base import BertEncoder, get_model_obj
  from blink.common.optimizer import get_bert_optimizer
  from elq.biencoder.allennlp_span_utils import batched_span_select, batched_index_select
  from elq.biencoder.utils import batch_reshape_mask_left
label_desc: Use image node_id 6 for calling the BiEncoderRanker obj's underlying member method code with example usage: obj.encode_context(cands, gold_mention_bounds, gold_mention_bounds_mask, num_cand_mentions, topK_threshold, get_mention_scores) and returns: context_outs_reshape, context_outs, unknown