code: 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
  
  
  _SCREAMING_SNAKE_CASE   						=    logging.get_logger(__name__)
  
  _SCREAMING_SNAKE_CASE   						=    [
      """yjernite/retribert-base-uncased""",
      # See all RetriBert models at https://huggingface.co/models?filter=retribert
  ]
  
  
  
  
  
  
  
  class 					UpperCAmelCase_				(							_a    ):
      UpperCAmelCase_			:       int						 =							RetriBertConfig
      UpperCAmelCase_			:       Optional[int]						 =							None
      UpperCAmelCase_			:       Any						 =							'retribert'
  
  
  
  
      def 						__lowerCamelCase (					self  :   List[Any]      ,      lowerCAmelCase  :   Tuple   )	-> str:
        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   )
  
  _SCREAMING_SNAKE_CASE   						=    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.'     ,     _a     ,     )
  class 					UpperCAmelCase_				(							_a    ):
  
  
  
  
      def __init__(					self  :   Dict      ,      lowerCAmelCase  :   RetriBertConfig   )	-> None:
        super().__init__(lowerCAmelCase   )
        a__     		=     config.projection_dim
  
        a__     		=     BertModel(lowerCAmelCase   )
        a__     		=     None if config.share_encoders else BertModel(lowerCAmelCase   )
        a__     		=     nn.Dropout(config.hidden_dropout_prob   )
        a__     		=     nn.Linear(config.hidden_size      ,      config.projection_dim      ,      bias=lowerCAmelCase   )
        a__     		=     nn.Linear(config.hidden_size      ,      config.projection_dim      ,      bias=lowerCAmelCase   )
  
        a__     		=     nn.CrossEntropyLoss(reduction="mean"   )
  
        # Initialize weights and apply final processing
        self.post_init()
  
  
  
  
      def 						__lowerCamelCase (					self  :   List[str]      ,      lowerCAmelCase  :   Dict      ,      lowerCAmelCase  :   Tuple      ,      lowerCAmelCase  :   Tuple      ,      lowerCAmelCase  :   Optional[int]=-1      ,      )	-> Dict:
        # reproduces BERT forward pass with checkpointing
        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
          a__     		=     input_ids.device
          a__     		=     input_ids.size()
          a__     		=     torch.zeros(lowerCAmelCase      ,      dtype=torch.long      ,      device=lowerCAmelCase   )
          a__     		=     [None] * sent_encoder.config.num_hidden_layers
          a__     		=     sent_encoder.get_extended_attention_mask(
              lowerCAmelCase      ,      lowerCAmelCase   )
  
          # define function for checkpointing
          def partial_encode(*lowerCAmelCase  :   Optional[Any]   ):
            a__     		=     sent_encoder.encoder(
                inputs[0]      ,      attention_mask=inputs[1]      ,      head_mask=lowerCAmelCase      ,      )
            a__     		=     encoder_outputs[0]
            a__     		=     sent_encoder.pooler(lowerCAmelCase   )
            return pooled_output
  
          # run embedding layer on everything at once
          a__     		=     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
          a__     		=     []
          for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size   )   ):
            a__     		=     embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
            a__     		=     extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
            a__     		=     checkpoint.checkpoint(lowerCAmelCase      ,      lowerCAmelCase      ,      lowerCAmelCase   )
            pooled_output_list.append(lowerCAmelCase   )
          return torch.cat(lowerCAmelCase      ,      dim=0   )
  
  
  
  
      def 						__lowerCamelCase (					self  :   List[Any]      ,      lowerCAmelCase  :   Dict      ,      lowerCAmelCase  :   List[Any]=None      ,      lowerCAmelCase  :   str=-1      ,      )	-> Dict:
        a__     		=     self.embed_sentences_checkpointed(
            lowerCAmelCase      ,      lowerCAmelCase      ,      self.bert_query      ,      lowerCAmelCase      ,      )
        return self.project_query(lowerCAmelCase   )
  
  
  
  
      def 						__lowerCamelCase (					self  :   Optional[Any]      ,      lowerCAmelCase  :   Optional[Any]      ,      lowerCAmelCase  :   Optional[Any]=None      ,      lowerCAmelCase  :   List[Any]=-1      ,      )	-> List[str]:
        a__     		=     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 						__lowerCamelCase (					self  :   int      ,      lowerCAmelCase  :   torch.LongTensor      ,      lowerCAmelCase  :   Optional[torch.FloatTensor]      ,      lowerCAmelCase  :   torch.LongTensor      ,      lowerCAmelCase  :   Optional[torch.FloatTensor]      ,      lowerCAmelCase  :   int = -1      ,      )	-> torch.FloatTensor:
        a__     		=     input_ids_query.device
        a__     		=     self.embed_questions(lowerCAmelCase      ,      lowerCAmelCase      ,      lowerCAmelCase   )
        a__     		=     self.embed_answers(lowerCAmelCase      ,      lowerCAmelCase      ,      lowerCAmelCase   )
        a__     		=     torch.mm(lowerCAmelCase      ,      a_reps.t()   )
        a__     		=     self.ce_loss(lowerCAmelCase      ,      torch.arange(compare_scores.shape[1]   ).to(lowerCAmelCase   )   )
        a__     		=     self.ce_loss(compare_scores.t()      ,      torch.arange(compare_scores.shape[0]   ).to(lowerCAmelCase   )   )
        a__     		=     (loss_qa + loss_aq) / 2
        return loss
  
  
style_context: def 				__lowercase	(						lowerCamelCase__ = "mytext.txt" ,				lowerCamelCase__ = False					)				->None:
  
    '''simple docstring'''
  
  
  
  
    import socket
  
    a__     		=     1_23_12  # Reserve a port for your service.
    a__     		=     socket.socket()  # Create a socket object
    a__     		=     socket.gethostname()  # Get local machine name
    sock.bind((host, port)					)  # Bind to the port
    sock.listen(5					)  # Now wait for client connection.
  
    print("Server listening...."					)
  
    while True:
      a__		,a__     		=     sock.accept()  # Establish connection with client.
      print(f"Got connection from {addr}"					)
      a__     		=     conn.recv(10_24					)
      print(f"Server received: {data = }"					)
  
      with open(lowerCamelCase__ ,				"rb"					) as in_file:
        a__     		=     in_file.read(10_24					)
        while data:
          conn.send(lowerCamelCase__					)
          print(f"Sent {data!r}"					)
          a__     		=     in_file.read(10_24					)
  
      print("Done sending"					)
      conn.close()
      if testing:  # Allow the test to complete
        break
  
    sock.shutdown(1					)
    sock.close()
  
  
  if __name__ == "__main__":
       send_file()
  
  