class_docstr: Bert-based module for Biencoder model.
class_docstr_tok: ["Bert-based","module","for","Biencoder","model","."]
code: def __init__(
      self,
      num_tokentypes=1,
      parallel_output=True,
      only_query_model=False,
      only_context_model=False,
      biencoder_shared_query_context_model=False,
      pre_process=True,
      post_process=True,
  ):
      super(BiEncoderModel, self).__init__()
      args = get_args()
  
      bert_kwargs = dict(
          num_tokentypes=num_tokentypes,
          parallel_output=parallel_output,
          pre_process=pre_process,
          post_process=post_process,
      )
  
      self.biencoder_shared_query_context_model = (
          biencoder_shared_query_context_model
      )
      assert not (only_context_model and only_query_model)
      self.use_context_model = not only_query_model
      self.use_query_model = not only_context_model
      self.biencoder_projection_dim = args.biencoder_projection_dim
  
      if self.biencoder_shared_query_context_model:
          self.model = PretrainedBertModel(**bert_kwargs)
          self._model_key = "shared_model"
          self.query_model, self.context_model = self.model, self.model
      else:
          if self.use_query_model:
              # this model embeds (pseudo-)queries - Embed_input in the paper
              self.query_model = PretrainedBertModel(**bert_kwargs)
              self._query_key = "query_model"
  
          if self.use_context_model:
              # this model embeds evidence blocks - Embed_doc in the paper
              self.context_model = PretrainedBertModel(**bert_kwargs)
              self._context_key = "context_model"
  
imports: import os
  import torch
  import sys
  from megatron import get_args, print_rank_0, get_tokenizer
  from megatron.core import mpu
  from megatron.checkpointing import fix_query_key_value_ordering
  from megatron.checkpointing import get_checkpoint_tracker_filename
  from megatron.checkpointing import get_checkpoint_name
  from megatron.model.bert_model import bert_position_ids
  from megatron.model.enums import AttnMaskType
  from megatron.model.language_model import get_language_model
  from megatron.model.utils import get_linear_layer
  from megatron.model.utils import init_method_normal
  from megatron.model.utils import scaled_init_method_normal
  from .module import MegatronModule
label_desc: Use image node_id 1 to create a new BiEncoderModel object from inherited base classes: MegatronModule with example: obj = BiEncoderModel(num_tokentypes, parallel_output, only_query_model, only_context_model, biencoder_shared_query_context_model, pre_process, post_process)