docstr:     Prunes mentions based on mention scores/logits (by either
      `threshold` or `num_cand_mentions`, whichever yields less candidates)
  
  Inputs:
      mention_logits: torch.FloatTensor (bsz, num_total_mentions)
      mention_bounds: torch.IntTensor (bsz, num_total_mentions)
      num_cand_mentions: int
      threshold: float
  Returns:
      torch.FloatTensor(bsz, max_num_pred_mentions): top mention scores/logits
      torch.IntTensor(bsz, max_num_pred_mentions, 2): top mention boundaries
      torch.BoolTensor(bsz, max_num_pred_mentions): mask on top mentions
      torch.BoolTensor(bsz, total_possible_mentions): mask for reshaping from total possible mentions -> max # pred mentions
docstr_tok: ["Prunes","mentions","based","on","mention","scores\/logits","(","by","either","`","threshold","`","or","`","num_cand_mentions","`",",","whichever","yields","less","candidates",")","Inputs",":","mention_logits",":","torch.FloatTensor","(","bsz",",","num_total_mentions",")","mention_bounds",":","torch.IntTensor","(","bsz",",","num_total_mentions",")","num_cand_mentions",":","int","threshold",":","float","Returns",":","torch.FloatTensor","(","bsz",",","max_num_pred_mentions",")",":","top","mention","scores\/logits","torch.IntTensor","(","bsz",",","max_num_pred_mentions",",","2",")",":","top","mention","boundaries","torch.BoolTensor","(","bsz",",","max_num_pred_mentions",")",":","mask","on","top","mentions","torch.BoolTensor","(","bsz",",","total_possible_mentions",")",":","mask","for","reshaping","from","total","possible","mentions","-",">","max","#","pred","mentions"]
code: def prune_ctxt_mentions(
      self,
      mention_logits,
      mention_bounds,
      num_cand_mentions,
      threshold,
  ):
      """
          Prunes mentions based on mention scores/logits (by either
          `threshold` or `num_cand_mentions`, whichever yields less candidates)
  
      Inputs:
          mention_logits: torch.FloatTensor (bsz, num_total_mentions)
          mention_bounds: torch.IntTensor (bsz, num_total_mentions)
          num_cand_mentions: int
          threshold: float
      Returns:
          torch.FloatTensor(bsz, max_num_pred_mentions): top mention scores/logits
          torch.IntTensor(bsz, max_num_pred_mentions, 2): top mention boundaries
          torch.BoolTensor(bsz, max_num_pred_mentions): mask on top mentions
          torch.BoolTensor(bsz, total_possible_mentions): mask for reshaping from total possible mentions -> max # pred mentions
      """
      # (bsz, num_cand_mentions); (bsz, num_cand_mentions)
      top_mention_logits, mention_pos = mention_logits.topk(
          num_cand_mentions, sorted=True
      )
      # (bsz, num_cand_mentions, 2)
      #   [:,:,0]: index of batch
      #   [:,:,1]: index into top mention in mention_bounds
      mention_pos = torch.stack(
          [
              torch.arange(mention_pos.size(0))
              .to(mention_pos.device)
              .unsqueeze(-1)
              .expand_as(mention_pos),
              mention_pos,
          ],
          dim=-1,
      )
      # (bsz, num_cand_mentions)
      top_mention_pos_mask = (
          torch.sigmoid(top_mention_logits).log() > threshold
      )
      # (total_possible_mentions, 2)
      #   tuples of [index of batch, index into mention_bounds] of what mentions to include
      mention_pos = mention_pos[
          top_mention_pos_mask
          | (
              # 2nd part of OR: if nothing is > threshold, use topK that are > -inf
              ((top_mention_pos_mask.sum(1) == 0).unsqueeze(-1))
              & (top_mention_logits > -float("inf"))
          )
      ]
      mention_pos = mention_pos.view(-1, 2)
      # (bsz, total_possible_mentions)
      #   mask of possible logits
      mention_pos_mask = torch.zeros(
          mention_logits.size(), dtype=torch.bool
      ).to(mention_pos.device)
      mention_pos_mask[mention_pos[:, 0], mention_pos[:, 1]] = 1
      # (bsz, max_num_pred_mentions, 2)
      (
          chosen_mention_bounds,
          chosen_mention_mask,
      ) = batch_reshape_mask_left(
          mention_bounds, mention_pos_mask, pad_idx=0
      )
      # (bsz, max_num_pred_mentions)
      chosen_mention_logits, _ = batch_reshape_mask_left(
          mention_logits,
          mention_pos_mask,
          pad_idx=-float("inf"),
          left_align_mask=chosen_mention_mask,
      )
      return (
          chosen_mention_logits,
          chosen_mention_bounds,
          chosen_mention_mask,
          mention_pos_mask,
      )
  
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 4 for calling the BiEncoderModule obj's underlying member method code with example usage: obj.prune_ctxt_mentions(mention_logits, mention_bounds, num_cand_mentions, threshold) and returns: chosen_mention_logits,  chosen_mention_bounds,  chosen_mention_mask,  mention_pos_mask