code: 
  
  """simple docstring"""
  import math
  from typing import Optional
  
  import torch
  import torch.utils.checkpoint as checkpoint
  from torch import nn
  
  from ....modeling_utils import PreTrainedModel
  from ....utils import add_start_docstrings, logging
  from ...bert.modeling_bert import BertModel
  from .configuration_retribert import RetriBertConfig
  
  
  _A						=     logging.get_logger(__name__)
  
  _A						=     [
      """yjernite/retribert-base-uncased""",
      # See all RetriBert models at https://huggingface.co/models?filter=retribert
  ]
  
  class 			lowerCamelCase  (					lowerCAmelCase__      ):
  
  
    '''simple docstring'''
  
  
  
  
  
  
  
    SCREAMING_SNAKE_CASE							=     RetriBertConfig
    SCREAMING_SNAKE_CASE							=     None
    SCREAMING_SNAKE_CASE							=     'retribert'
  
  
  
    def    _a			(self       ,					_lowerCamelCase  ):
  
  
          """simple docstring"""
  
          if isinstance(_lowerCamelCase       ,					nn.Linear  ):
                module.weight.data.normal_(mean=0.0       ,					std=self.config.initializer_range  )
                if module.bias is not None:
                      module.bias.data.zero_()
          elif isinstance(_lowerCamelCase       ,					nn.Embedding  ):
                module.weight.data.normal_(mean=0.0       ,					std=self.config.initializer_range  )
                if module.padding_idx is not None:
                      module.weight.data[module.padding_idx].zero_()
          elif isinstance(_lowerCamelCase       ,					nn.LayerNorm  ):
                module.bias.data.zero_()
                module.weight.data.fill_(1.0  )
  
  
  
  
  _A						=     R"""
  
      This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
      library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
      etc.)
  
      This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
      Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
      and behavior.
  
      Parameters:
          config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.
              Initializing with a config file does not load the weights associated with the model, only the
              configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
  """
  
  @add_start_docstrings(
      'Bert Based model to embed queries or document for document retrieval.'				,				lowerCAmelCase__				,				)
  class 			lowerCamelCase  (					lowerCAmelCase__      ):
  
  
    '''simple docstring'''
  
  
  
  
  
  
  
    def __init__(self       ,					_lowerCamelCase  ):
  
  
          """simple docstring"""
  
          super().__init__(_lowerCamelCase  )
          UpperCAmelCase__ :       Dict						=						config.projection_dim
  
          UpperCAmelCase__ :       Optional[Any]						=						BertModel(_lowerCamelCase  )
          UpperCAmelCase__ :       Optional[int]						=						None if config.share_encoders else BertModel(_lowerCamelCase  )
          UpperCAmelCase__ :       Optional[int]						=						nn.Dropout(config.hidden_dropout_prob  )
          UpperCAmelCase__ :       Dict						=						nn.Linear(config.hidden_size       ,					config.projection_dim       ,					bias=_lowerCamelCase  )
          UpperCAmelCase__ :       Tuple						=						nn.Linear(config.hidden_size       ,					config.projection_dim       ,					bias=_lowerCamelCase  )
  
          UpperCAmelCase__ :       Optional[Any]						=						nn.CrossEntropyLoss(reduction="""mean"""  )
  
          # Initialize weights and apply final processing
          self.post_init()
  
  
  
    def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase=-1       ,					):
  
  
          """simple docstring"""
  
          if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
                return sent_encoder(_lowerCamelCase       ,					attention_mask=_lowerCamelCase  )[1]
          else:
                # prepare implicit variables
                UpperCAmelCase__ :       str						=						input_ids.device
                UpperCAmelCase__ :       str						=						input_ids.size()
                UpperCAmelCase__ :       List[str]						=						torch.zeros(_lowerCamelCase       ,					dtype=torch.long       ,					device=_lowerCamelCase  )
                UpperCAmelCase__ :       List[str]						=						[None] * sent_encoder.config.num_hidden_layers
                UpperCAmelCase__ :       torch.Tensor						=						sent_encoder.get_extended_attention_mask(
                    _lowerCamelCase       ,					_lowerCamelCase  )
  
                # define function for checkpointing
                def partial_encode(*_lowerCamelCase  ):
                      UpperCAmelCase__ :       Optional[int]						=						sent_encoder.encoder(
                          inputs[0]       ,					attention_mask=inputs[1]       ,					head_mask=_lowerCamelCase       ,					)
                      UpperCAmelCase__ :       Any						=						encoder_outputs[0]
                      UpperCAmelCase__ :       int						=						sent_encoder.pooler(_lowerCamelCase  )
                      return pooled_output
  
                # run embedding layer on everything at once
                UpperCAmelCase__ :       Optional[Any]						=						sent_encoder.embeddings(
                    input_ids=_lowerCamelCase       ,					position_ids=_lowerCamelCase       ,					token_type_ids=_lowerCamelCase       ,					inputs_embeds=_lowerCamelCase  )
                # run encoding and pooling on one mini-batch at a time
                UpperCAmelCase__ :       Union[str, Any]						=						[]
                for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size  )  ):
                      UpperCAmelCase__ :       List[str]						=						embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                      UpperCAmelCase__ :       Dict						=						extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
                      UpperCAmelCase__ :       List[str]						=						checkpoint.checkpoint(_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  )
                      pooled_output_list.append(_lowerCamelCase  )
                return torch.cat(_lowerCamelCase       ,					dim=0  )
  
  
  
    def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase=None       ,					_lowerCamelCase=-1       ,					):
  
  
          """simple docstring"""
  
          UpperCAmelCase__ :       Dict						=						self.embed_sentences_checkpointed(
              _lowerCamelCase       ,					_lowerCamelCase       ,					self.bert_query       ,					_lowerCamelCase       ,					)
          return self.project_query(_lowerCamelCase  )
  
  
  
    def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase=None       ,					_lowerCamelCase=-1       ,					):
  
  
          """simple docstring"""
  
          UpperCAmelCase__ :       Optional[int]						=						self.embed_sentences_checkpointed(
              _lowerCamelCase       ,					_lowerCamelCase       ,					self.bert_query if self.bert_doc is None else self.bert_doc       ,					_lowerCamelCase       ,					)
          return self.project_doc(_lowerCamelCase  )
  
  
  
  
    def    _a			(self       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase = -1       ,					):
  
  
          """simple docstring"""
  
          UpperCAmelCase__ :       Optional[Any]						=						input_ids_query.device
          UpperCAmelCase__ :       str						=						self.embed_questions(_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  )
          UpperCAmelCase__ :       Union[str, Any]						=						self.embed_answers(_lowerCamelCase       ,					_lowerCamelCase       ,					_lowerCamelCase  )
          UpperCAmelCase__ :       Any						=						torch.mm(_lowerCamelCase       ,					a_reps.t()  )
          UpperCAmelCase__ :       Any						=						self.ce_loss(_lowerCamelCase       ,					torch.arange(compare_scores.shape[1]  ).to(_lowerCamelCase  )  )
          UpperCAmelCase__ :       List[str]						=						self.ce_loss(compare_scores.t()       ,					torch.arange(compare_scores.shape[0]  ).to(_lowerCamelCase  )  )
          UpperCAmelCase__ :       str						=						(loss_qa + loss_aq) / 2
          return loss
  
  
style_context: 
  
  """simple docstring"""
  from typing import TYPE_CHECKING
  
  from ...utils import (
      OptionalDependencyNotAvailable,
      _LazyModule,
      is_flax_available,
      is_keras_nlp_available,
      is_tensorflow_text_available,
      is_tf_available,
      is_tokenizers_available,
      is_torch_available,
  )
  
  
  _A						=     {
      """configuration_gpt2""": ["""GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPT2Config""", """GPT2OnnxConfig"""],
      """tokenization_gpt2""": ["""GPT2Tokenizer"""],
  }
  
  try:
      if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
      pass
  else:
      _A						=     ["""GPT2TokenizerFast"""]
  
  try:
      if not is_torch_available():
          raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
      pass
  else:
      _A						=     [
          """GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """GPT2DoubleHeadsModel""",
          """GPT2ForQuestionAnswering""",
          """GPT2ForSequenceClassification""",
          """GPT2ForTokenClassification""",
          """GPT2LMHeadModel""",
          """GPT2Model""",
          """GPT2PreTrainedModel""",
          """load_tf_weights_in_gpt2""",
      ]
  
  try:
      if not is_tf_available():
          raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
      pass
  else:
      _A						=     [
          """TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """TFGPT2DoubleHeadsModel""",
          """TFGPT2ForSequenceClassification""",
          """TFGPT2LMHeadModel""",
          """TFGPT2MainLayer""",
          """TFGPT2Model""",
          """TFGPT2PreTrainedModel""",
      ]
  
  try:
      if not is_keras_nlp_available():
          raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
      pass
  else:
      _A						=     ["""TFGPT2Tokenizer"""]
  
  try:
      if not is_flax_available():
          raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
      pass
  else:
      _A						=     ["""FlaxGPT2LMHeadModel""", """FlaxGPT2Model""", """FlaxGPT2PreTrainedModel"""]
  
  if TYPE_CHECKING:
      from .configuration_gpta import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTaConfig, GPTaOnnxConfig
      from .tokenization_gpta import GPTaTokenizer
  
      try:
          if not is_tokenizers_available():
              raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
          pass
      else:
          from .tokenization_gpta_fast import GPTaTokenizerFast
  
      try:
          if not is_torch_available():
              raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
          pass
      else:
          from .modeling_gpta import (
              GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTaDoubleHeadsModel,
              GPTaForQuestionAnswering,
              GPTaForSequenceClassification,
              GPTaForTokenClassification,
              GPTaLMHeadModel,
              GPTaModel,
              GPTaPreTrainedModel,
              load_tf_weights_in_gpta,
          )
  
      try:
          if not is_tf_available():
              raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
          pass
      else:
          from .modeling_tf_gpta import (
              TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFGPTaDoubleHeadsModel,
              TFGPTaForSequenceClassification,
              TFGPTaLMHeadModel,
              TFGPTaMainLayer,
              TFGPTaModel,
              TFGPTaPreTrainedModel,
          )
  
      try:
          if not is_keras_nlp_available():
              raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
          pass
      else:
          from .tokenization_gpta_tf import TFGPTaTokenizer
  
      try:
          if not is_flax_available():
              raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
          pass
      else:
          from .modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel, FlaxGPTaPreTrainedModel
  
  else:
      import sys
  
      _A						=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
  
  