


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


UpperCAmelCase__     			=							logging.get_logger(__name__)

UpperCAmelCase__     			=							[
    '''yjernite/retribert-base-uncased''',
    # See all RetriBert models at https://huggingface.co/models?filter=retribert
]







class    lowerCamelCase__     ( lowerCAmelCase):
     SCREAMING_SNAKE_CASE__       =    RetriBertConfig
     SCREAMING_SNAKE_CASE__       =    None
     SCREAMING_SNAKE_CASE__       =    '''retribert'''





     def   __A      (self	,  UpperCAmelCase )  ->     Optional[Any]:
      if isinstance(UpperCAmelCase	,  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(UpperCAmelCase	,  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(UpperCAmelCase	,  nn.LayerNorm ):
       module.bias.data.zero_()
       module.weight.data.fill_(1.0 )



UpperCAmelCase__     			=							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):





     def __init__(self	,  UpperCAmelCase )  ->     None:
      super().__init__(UpperCAmelCase )
      _lowercase							   =config.projection_dim

      _lowercase							   =BertModel(UpperCAmelCase )
      _lowercase							   =None if config.share_encoders else BertModel(UpperCAmelCase )
      _lowercase							   =nn.Dropout(config.hidden_dropout_prob )
      _lowercase							   =nn.Linear(config.hidden_size	,  config.projection_dim	,  bias=UpperCAmelCase )
      _lowercase							   =nn.Linear(config.hidden_size	,  config.projection_dim	,  bias=UpperCAmelCase )

      _lowercase							   =nn.CrossEntropyLoss(reduction='''mean''' )

      # Initialize weights and apply final processing
      self.post_init()





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase=-1	,  )  ->     Tuple:
      # reproduces BERT forward pass with checkpointing
      if checkpoint_batch_size < 0 or input_ids.shape[0] < checkpoint_batch_size:
       return sent_encoder(UpperCAmelCase	,  attention_mask=UpperCAmelCase )[1]
      else:
       # prepare implicit variables
       _lowercase							   =input_ids.device
       _lowercase							   =input_ids.size()
       _lowercase							   =torch.zeros(UpperCAmelCase	,  dtype=torch.long	,  device=UpperCAmelCase )
       _lowercase							   =[None] * sent_encoder.config.num_hidden_layers
       _lowercase							   =sent_encoder.get_extended_attention_mask(
           UpperCAmelCase	,  UpperCAmelCase )

       # define function for checkpointing
       def partial_encode(*UpperCAmelCase ):
        _lowercase							   =sent_encoder.encoder(
            inputs[0]	,  attention_mask=inputs[1]	,  head_mask=UpperCAmelCase	,  )
        _lowercase							   =encoder_outputs[0]
        _lowercase							   =sent_encoder.pooler(UpperCAmelCase )
        return pooled_output

       # run embedding layer on everything at once
       _lowercase							   =sent_encoder.embeddings(
           input_ids=UpperCAmelCase	,  position_ids=UpperCAmelCase	,  token_type_ids=UpperCAmelCase	,  inputs_embeds=UpperCAmelCase )
       # run encoding and pooling on one mini-batch at a time
       _lowercase							   =[]
       for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size ) ):
        _lowercase							   =embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
        _lowercase							   =extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
        _lowercase							   =checkpoint.checkpoint(UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )
        pooled_output_list.append(UpperCAmelCase )
       return torch.cat(UpperCAmelCase	,  dim=0 )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase=None	,  UpperCAmelCase=-1	,  )  ->     Union[str, Any]:
      _lowercase							   =self.embed_sentences_checkpointed(
          UpperCAmelCase	,  UpperCAmelCase	,  self.bert_query	,  UpperCAmelCase	,  )
      return self.project_query(UpperCAmelCase )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase=None	,  UpperCAmelCase=-1	,  )  ->     Optional[Any]:
      _lowercase							   =self.embed_sentences_checkpointed(
          UpperCAmelCase	,  UpperCAmelCase	,  self.bert_query if self.bert_doc is None else self.bert_doc	,  UpperCAmelCase	,  )
      return self.project_doc(UpperCAmelCase )





     def   __A      (self	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase = -1	,  )  ->     torch.FloatTensor:
      _lowercase							   =input_ids_query.device
      _lowercase							   =self.embed_questions(UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )
      _lowercase							   =self.embed_answers(UpperCAmelCase	,  UpperCAmelCase	,  UpperCAmelCase )
      _lowercase							   =torch.mm(UpperCAmelCase	,  a_reps.t() )
      _lowercase							   =self.ce_loss(UpperCAmelCase	,  torch.arange(compare_scores.shape[1] ).to(UpperCAmelCase ) )
      _lowercase							   =self.ce_loss(compare_scores.t()	,  torch.arange(compare_scores.shape[0] ).to(UpperCAmelCase ) )
      _lowercase							   =(loss_qa + loss_aq) / 2
      return loss





import unittest
from functools import partial

import torch

from accelerate import Accelerator, debug_launcher
from accelerate.state import AcceleratorState, GradientState
from accelerate.test_utils import require_cpu, require_huggingface_suite
from accelerate.utils import GradientAccumulationPlugin
def        UpperCAmelCase_      (							__snake_case=2  ,		__snake_case=True  ,		__snake_case=False		)  ->    Dict:







 """simple docstring"""


 _lowercase							   =Accelerator(step_scheduler_with_optimizer=__snake_case  ,		split_batches=__snake_case		)
 _lowercase							   =torch.nn.Linear(2  ,		4		)
 _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=1.0		)
 _lowercase							   =torch.optim.lr_scheduler.OneCycleLR(__snake_case  ,		max_lr=0.01  ,		steps_per_epoch=2  ,		epochs=1		)
 _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

 # Optimizer has stepped
 scheduler.step()
 if step_scheduler_with_optimizer or (num_processes == 1):
  assert (
      scheduler.scheduler.last_epoch == num_processes
  ), F"Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})"
 else:
  assert (
      scheduler.scheduler.last_epoch != num_processes
  ), F"Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})"
def        UpperCAmelCase_      (							__snake_case=2  ,		__snake_case=True  ,		__snake_case=False		)  ->    Optional[int]:







 """simple docstring"""


 _lowercase							   =Accelerator(step_scheduler_with_optimizer=__snake_case  ,		split_batches=__snake_case		)
 _lowercase							   =torch.nn.Linear(2  ,		4		)
 _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=1.0		)
 _lowercase							   =torch.optim.lr_scheduler.LambdaLR(__snake_case  ,		lr_lambda=lambda __snake_case		: 1 - n / 10		)
 _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

 # Optimizer has stepped
 _lowercase							   =False
 scheduler.step()
 _lowercase							   =1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F"Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"

 # Optimizer has not stepped
 _lowercase							   =True
 scheduler.step()
 if not step_scheduler_with_optimizer:
  _lowercase							   =1 - 2 / 10
 assert (
     scheduler.get_last_lr()[0] == expected_lr
 ), F"Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}"
def        UpperCAmelCase_      (							__snake_case = 2		)  ->    Optional[int]:







 """simple docstring"""


 from transformers import get_linear_schedule_with_warmup

 _lowercase							   =[1, 2, 4]
 for num_steps in steps:
  _lowercase							   =GradientAccumulationPlugin(num_steps=__snake_case  ,		adjust_scheduler=num_steps > 1		)
  _lowercase							   =Accelerator(gradient_accumulation_plugin=__snake_case		)
  _lowercase							   =torch.nn.Linear(2  ,		4		)
  _lowercase							   =torch.optim.AdamW(model.parameters()  ,		lr=10.0		)
  _lowercase							   =get_linear_schedule_with_warmup(optimizer=__snake_case  ,		num_warmup_steps=0  ,		num_training_steps=20		)

  _lowercase   ,					_lowercase   ,					_lowercase							   =accelerator.prepare(__snake_case  ,		__snake_case  ,		__snake_case		)

  for i in range(10 * num_steps		):
   with accelerator.accumulate(__snake_case		):
    optimizer.step()
    scheduler.step()

   if i == (10 * num_steps - 2):
    assert (
        scheduler.get_last_lr()[0] != 0
    ), F"Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}"
  assert (
      scheduler.get_last_lr()[0] == 0
  ), F"Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}"
  GradientState._reset_state()







@require_cpu
class    lowerCamelCase__     ( unittest.TestCase):





     def   __A      (self )  ->     Optional[Any]:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 )	,  num_processes=1 )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Any:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 )	,  num_processes=1 )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Union[str, Any]:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  step_scheduler_with_optimizer=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Dict:
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  step_scheduler_with_optimizer=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     int:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     Optional[Any]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1	,  split_batches=UpperCAmelCase )	,  num_processes=1 )





     def   __A      (self )  ->     List[str]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  step_scheduler_with_optimizer=UpperCAmelCase ) )





     def   __A      (self )  ->     Union[str, Any]:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  step_scheduler_with_optimizer=UpperCAmelCase ) )





     @require_huggingface_suite
     def   __A      (self )  ->     int:
      AcceleratorState._reset_state(UpperCAmelCase )
      debug_launcher(partial(UpperCAmelCase	,  num_processes=1 ) )
      debug_launcher(UpperCAmelCase )

