code: def __init__(self, params):
      super(BiEncoderModule, self).__init__()
      ctxt_bert = BertModel.from_pretrained(
          params["bert_model"], output_hidden_states=True
      )
      if params["load_cand_enc_only"]:
          bert_model = "bert-large-uncased"
      else:
          bert_model = params["bert_model"]
      cand_bert = BertModel.from_pretrained(
          bert_model,
          output_hidden_states=True,
      )
      self.context_encoder = BertEncoder(
          ctxt_bert,
          params["out_dim"],
          layer_pulled=params["pull_from_layer"],
          add_linear=params["add_linear"],
      )
      self.cand_encoder = BertEncoder(
          cand_bert,
          params["out_dim"],
          layer_pulled=params["pull_from_layer"],
          add_linear=params["add_linear"],
      )
      if params.get("freeze_cand_enc", False):
          for param in self.cand_encoder.parameters():
              param.requires_grad = False
  
      self.config = ctxt_bert.config
  
      ctxt_bert_output_dim = (
          ctxt_bert.embeddings.word_embeddings.weight.size(1)
      )
  
      self.mention_aggregation_type = params.get(
          "mention_aggregation_type", None
      )
      self.classification_heads = nn.ModuleDict({})
      self.linear_compression = None
      if self.mention_aggregation_type is not None:
          classification_heads_dict = {
              "get_context_embeds": GetContextEmbedsHead(
                  self.mention_aggregation_type,
                  ctxt_bert_output_dim,
                  cand_bert.embeddings.word_embeddings.weight.size(1),
              )
          }
          classification_heads_dict[
              "mention_scores"
          ] = MentionScoresHead(
              ctxt_bert_output_dim,
              params["mention_scoring_method"],
              params.get("max_mention_length", 10),
          )
          self.classification_heads = nn.ModuleDict(
              classification_heads_dict
          )
      elif (
          ctxt_bert_output_dim
          != cand_bert.embeddings.word_embeddings.weight.size(1)
      ):
          # mapping to make the output dimensions match for dot-product similarity
          self.linear_compression = nn.Linear(
              ctxt_bert_output_dim,
              cand_bert.embeddings.word_embeddings.weight.size(1),
          )
  
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 1 to create a new BiEncoderModule object from inherited base classes: torch.nn with example: obj = BiEncoderModule(params)