


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		:Tuple											=       logging.get_logger(__name__)

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


class       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  _SCREAMING_SNAKE_CASE						      =    RetriBertConfig
  _SCREAMING_SNAKE_CASE						      =    None
  _SCREAMING_SNAKE_CASE						      =    'retribert'



  def    lowerCAmelCase__					(       self    :    str      ,			_lowerCAmelCase    :    List[str]					)				->	Any:




     """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					)






SCREAMING_SNAKE_CASE		:str											=       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       __lowerCAmelCase       (	a   ):






  """simple docstring"""





  def __init__(       self    :    Dict      ,			_lowerCAmelCase    :    RetriBertConfig					)				->	None:




     """simple docstring"""
     super().__init__(_lowerCAmelCase					)
     snake_case_         =   config.projection_dim

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

     snake_case_         =   nn.CrossEntropyLoss(reduction="mean"					)

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



  def    lowerCAmelCase__					(       self    :    Any      ,			_lowerCAmelCase    :    str      ,			_lowerCAmelCase    :    List[str]      ,			_lowerCAmelCase    :    Tuple      ,			_lowerCAmelCase    :    str=-1      ,			)				->	str:




     """simple docstring"""
     # 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
        snake_case_         =   input_ids.device
        snake_case_         =   input_ids.size()
        snake_case_         =   torch.zeros(_lowerCAmelCase      ,			dtype=torch.long      ,			device=_lowerCAmelCase					)
        snake_case_         =   [None] * sent_encoder.config.num_hidden_layers
        snake_case_         =   sent_encoder.get_extended_attention_mask(
            _lowerCAmelCase      ,			_lowerCAmelCase					)

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

        # run embedding layer on everything at once
        snake_case_         =   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
        snake_case_         =   []
        for b in range(math.ceil(input_ids.shape[0] / checkpoint_batch_size					)					):
           snake_case_         =   embedding_output[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
           snake_case_         =   extended_attention_mask[b * checkpoint_batch_size : (b + 1) * checkpoint_batch_size]
           snake_case_         =   checkpoint.checkpoint(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)
           pooled_output_list.append(_lowerCAmelCase					)
        return torch.cat(_lowerCAmelCase      ,			dim=0					)



  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    Dict      ,			_lowerCAmelCase    :    List[str]=None      ,			_lowerCAmelCase    :    Union[str, Any]=-1      ,			)				->	int:




     """simple docstring"""
     snake_case_         =   self.embed_sentences_checkpointed(
         _lowerCAmelCase      ,			_lowerCAmelCase      ,			self.bert_query      ,			_lowerCAmelCase      ,			)
     return self.project_query(_lowerCAmelCase					)



  def    lowerCAmelCase__					(       self    :    int      ,			_lowerCAmelCase    :    int      ,			_lowerCAmelCase    :    List[str]=None      ,			_lowerCAmelCase    :    Union[str, Any]=-1      ,			)				->	Optional[Any]:




     """simple docstring"""
     snake_case_         =   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    :    List[Any]      ,			_lowerCAmelCase    :    torch.LongTensor      ,			_lowerCAmelCase    :    Optional[torch.FloatTensor]      ,			_lowerCAmelCase    :    torch.LongTensor      ,			_lowerCAmelCase    :    Optional[torch.FloatTensor]      ,			_lowerCAmelCase    :    int = -1      ,			)				->	torch.FloatTensor:




     """simple docstring"""
     snake_case_         =   input_ids_query.device
     snake_case_         =   self.embed_questions(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)
     snake_case_         =   self.embed_answers(_lowerCAmelCase      ,			_lowerCAmelCase      ,			_lowerCAmelCase					)
     snake_case_         =   torch.mm(_lowerCAmelCase      ,			a_reps.t()					)
     snake_case_         =   self.ce_loss(_lowerCAmelCase      ,			torch.arange(compare_scores.shape[1]					).to(_lowerCAmelCase					)					)
     snake_case_         =   self.ce_loss(compare_scores.t()      ,			torch.arange(compare_scores.shape[0]					).to(_lowerCAmelCase					)					)
     snake_case_         =   (loss_qa + loss_aq) / 2
     return loss








import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


SCREAMING_SNAKE_CASE		:Dict											=       logging.getLogger(__name__)



def      _lowerCAmelCase  (							lowerCAmelCase_							:List[str]							)->str:
   '''simple docstring'''






   random.seed(lowerCAmelCase_							)
   np.random.seed(lowerCAmelCase_							)
   torch.manual_seed(lowerCAmelCase_							)
   torch.cuda.manual_seed_all(lowerCAmelCase_							)



def      _lowerCAmelCase  (							lowerCAmelCase_							:Optional[int]						,  lowerCAmelCase_							:Union[str, Any]						,  lowerCAmelCase_							:Optional[Any]							)->Dict:
   '''simple docstring'''






   model.eval()
   snake_case_         =   next(model.parameters()							).device
   snake_case_         =   1
   snake_case_         =   torch.zeros((eval_batch_size, context_len)						,  dtype=torch.long						,  device=lowerCAmelCase_							)
   snake_case_         =   DataLoader(lowerCAmelCase_						,  shuffle=lowerCAmelCase_						,  batch_size=lowerCAmelCase_							)
   snake_case_         =   torch.zeros(1						,  device=lowerCAmelCase_							)
   snake_case_         =   0
   for batch in eval_dataloader:
      batch.to(lowerCAmelCase_							)
      # pad
      context.zero_()
      for i in range(lowerCAmelCase_							):
         snake_case_         =   batch[i]
      snake_case_         =   model(lowerCAmelCase_						,  labels=lowerCAmelCase_							)
      eval_loss += outputs[0].sum().item()
      nb_eval_examples += batch.size(0							)
   snake_case_         =   eval_loss / nb_eval_examples
   snake_case_         =   torch.exp(lowerCAmelCase_							)
   model.train()
   return perplexity



def      _lowerCAmelCase  (							lowerCAmelCase_							:Union[str, Any]="gpt2"							)->Union[str, Any]:
   '''simple docstring'''






   snake_case_         =   GPTaLMHeadModel.from_pretrained(lowerCAmelCase_						,  output_hidden_states=lowerCAmelCase_							)
   torch.save(model.state_dict()						,  model_name + "local.pt"							)
   return model



def      _lowerCAmelCase  (							lowerCAmelCase_							:List[str]						,  lowerCAmelCase_							:List[str]						,  lowerCAmelCase_							:List[Any]							)->int:
   '''simple docstring'''






   snake_case_         =   copy.deepcopy(lowerCAmelCase_							)
   model.to(lowerCAmelCase_							)

   snake_case_         =   ["bias", "LayerNorm.weight"]
   snake_case_         =   [
       {
           "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay							)],
           "weight_decay": 0.0,
       },
       {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay							)], "weight_decay": 0.0},
   ]
   snake_case_         =   AdamW(lowerCAmelCase_						,  lr=5e-5						,  eps=1e-8							)
   snake_case_         =   get_linear_schedule_with_warmup(lowerCAmelCase_						,  0						,  lowerCAmelCase_							)
   torch.cuda.empty_cache()
   return model, lm_optimizer, lm_scheduler



def      _lowerCAmelCase  (							lowerCAmelCase_							:Tuple						,  lowerCAmelCase_							:List[str]						,  lowerCAmelCase_							:Optional[int]						,  lowerCAmelCase_							:str							)->Tuple:
   '''simple docstring'''






   snake_case_         =   np.array(lowerCAmelCase_							).mean()
   snake_case_         =   np.array(lowerCAmelCase_							).std()
   snake_case_         =   (real_perps - avg) / std
   snake_case_         =   list(zip(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)							)
   joblib.dump(lowerCAmelCase_						,  lowerCAmelCase_							)



def      _lowerCAmelCase  (							lowerCAmelCase_							:Optional[int]						,  lowerCAmelCase_							:List[str]						,  lowerCAmelCase_							:Optional[Any]						,  lowerCAmelCase_							:str						,  lowerCAmelCase_							:Optional[int]						,  lowerCAmelCase_							:Any						,  lowerCAmelCase_							:Any						,  lowerCAmelCase_							:str="dev.jbl"						,  lowerCAmelCase_							:int=recopy_gpta						,  )->Tuple:
   '''simple docstring'''






   snake_case_         =   []
   snake_case_         =   []
   snake_case_         =   []

   # Initialize the transformer model
   snake_case_         =   copy.deepcopy(lowerCAmelCase_							)
   orig_model.to(device="cpu"							)
   torch.cuda.empty_cache()

   # Compute perplexity of initial transformer model for comparison
   model.train()
   snake_case_      ,      snake_case_      ,      snake_case_         =   recopy_model(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)

   for step in tqdm(range(lowerCAmelCase_							)							):
      snake_case_         =   torch.zeros((1, context_len)						,  dtype=torch.long						,  device=lowerCAmelCase_							)
      snake_case_         =   random.choice(lowerCAmelCase_							)
      snake_case_         =   random.randint(0						,  len(story[0]							) - context_len - 1							)
      snake_case_         =   story[0][start : start + context_len]
      lm_optimizer.zero_grad()
      snake_case_         =   model(lowerCAmelCase_						,  labels=lowerCAmelCase_							)
      snake_case_         =   outputs[0]
      snake_case_         =   compute_perplexity(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)
      model.train()
      lm_loss.backward()
      # Do LM backprop
      torch.nn.utils.clip_grad_norm_(model.parameters()						,  3.0							)
      lm_optimizer.step()
      lm_scheduler.step()  # Update learning rate schedule

      # Compute perplexity after back-propagating on the selected context
      snake_case_         =   compute_perplexity(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)

      # Periodically save the stored (X, IG(X)) pairs
      if step % 1_000 == 0 and step > 1:
         intermittent_save(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)

      # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
      snake_case_      ,      snake_case_      ,      snake_case_         =   recopy_model(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)

      past_perps.append(past_perp.item()							)
      real_perps.append(orig_perp - real_perp.item()							)
      contexts.append(np.array(context.cpu()							)							)

   intermittent_save(lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_						,  lowerCAmelCase_							)



def      _lowerCAmelCase  (							lowerCAmelCase_							:List[Any]						,  lowerCAmelCase_							:Union[str, Any]="data/tokenized_stories_train_wikitext103.jbl"						,  lowerCAmelCase_							:Tuple=100						,  lowerCAmelCase_							:int=1_026						,  lowerCAmelCase_							:int=True							)->List[Any]:
   '''simple docstring'''






   snake_case_         =   joblib.load(lowerCAmelCase_							)
   print("data loaded"							)
   snake_case_         =   []
   if trim:
      for i, example in enumerate(lowerCAmelCase_							):
         if len(example[0]							) > min_len:
            snake_case_         =   random.randint(0						,  len(example[0]							) - context_len - 1							)
            objective_set.append(example[0, start : start + context_len]							)
         if len(lowerCAmelCase_							) >= number:
            break
      snake_case_         =   []
      for j in range(i + 1						,  len(lowerCAmelCase_							)							):
         if len(data[j][0]							) > min_len:
            train_data.append(data[j]							)
   else:
      snake_case_         =   data[0:number]
      snake_case_         =   data[number:]

   joblib.dump(lowerCAmelCase_						,  "objective_set.jbl"							)
   print("objective set saved"							)
   return train_data, objective_set



def      _lowerCAmelCase  (							lowerCAmelCase_							:Union[str, Any]						,  lowerCAmelCase_							:Tuple						,  lowerCAmelCase_							:List[Any]						,  lowerCAmelCase_							:Optional[Any]						,  lowerCAmelCase_							:Dict=50						,  lowerCAmelCase_							:Dict="secondary_learner.pt"							)->Union[str, Any]:
   '''simple docstring'''






   snake_case_         =   torch.device("cuda:0" if torch.cuda.is_available() else "cpu"							)
   # We will use the first 512 pairs from our dataset as a test set for
   # our secondary learner and the rest to train
   snake_case_         =   train_dataset[:512]
   snake_case_         =   train_dataset[512:]
   snake_case_         =   DataLoader(lowerCAmelCase_						,  shuffle=lowerCAmelCase_						,  batch_size=lowerCAmelCase_							)
   snake_case_         =   DataLoader(lowerCAmelCase_						,  shuffle=lowerCAmelCase_						,  batch_size=lowerCAmelCase_							)

   # secondary learner model set up
   snake_case_         =   nn.MSELoss()
   snake_case_         =   nn.MSELoss(reduction="sum"							)
   secondary_learner.to(lowerCAmelCase_							)
   snake_case_         =   torch.optim.Adam(secondary_learner.parameters()						,  lr=0.0_0_0_0_1							)
   secondary_learner.train()

   # TODO in original code this is written as number of actual batches seen
   # not number of items seen but other places it is number of items instead.
   # improve consistency! changed this to epochs for clarity
   snake_case_         =   float("inf"							)
   # Iterate through batches until we've used max_steps batches
   for epoch in range(int(lowerCAmelCase_							)							):
      snake_case_         =   0.0
      secondary_learner.train()
      for step, batch in enumerate(lowerCAmelCase_							):
         snake_case_         =   batch[0].to(lowerCAmelCase_							)
         snake_case_         =   batch[1].to(lowerCAmelCase_							)
         snake_case_         =   secondary_learner(lowerCAmelCase_							)
         q_optimizer.zero_grad()
         snake_case_         =   loss(lowerCAmelCase_						,  real_q.float()							)
         q_loss.backward()
         q_optimizer.step()
         tr_q_loss += q_loss.item()

         # model trains fairly quickly so we won't wait for a full epoch
         # eval is triggered at eval_freq and end of epochs
         if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(lowerCAmelCase_							)):
            snake_case_         =   tr_q_loss / (step + 1)

            secondary_learner.eval()
            snake_case_         =   0.0
            snake_case_         =   0.0
            snake_case_         =   []
            snake_case_         =   []
            # Compute performance of the secondary learner after this batch
            for stepa, batcha in enumerate(lowerCAmelCase_							):
               snake_case_         =   batcha[0].to(lowerCAmelCase_							)
               snake_case_         =   batcha[1].to(lowerCAmelCase_							)
               snake_case_         =   secondary_learner(lowerCAmelCase_							)
               q_lossa += test_loss(lowerCAmelCase_						,  lowerCAmelCase_							).item()
               sum_qa += torch.sum(lowerCAmelCase_							).item()
               for ei, i in enumerate(predicted_qa.cpu().detach().numpy()							):
                  predicted.append(i.item()							)
               for ei, i in enumerate(real_qa.cpu().detach().numpy()							):
                  actual.append(i.item()							)

            q_lossa /= len(lowerCAmelCase_							)
            print(
                "Epoch: "						,  lowerCAmelCase_						,  "step: "						,  lowerCAmelCase_						,  "Avg. q:"						,  sum_qa / len(lowerCAmelCase_							)						,  "Train Loss: "						,  lowerCAmelCase_						,  "Test Loss: "						,  lowerCAmelCase_						,  )
            if q_lossa < best_test_loss:
               joblib.dump((predicted, actual)						,  "pred_vs_actual.jbl"							)
               torch.save(secondary_learner.state_dict()						,  lowerCAmelCase_							)
               snake_case_         =   q_lossa

         secondary_learner.train()
   return secondary_learner



class       __lowerCAmelCase       (	nn.Module   ):






  """simple docstring"""





  def __init__(       self    :    Optional[int]      ,			_lowerCAmelCase    :    Dict					)				->	List[Any]:




     """simple docstring"""
     super(_lowerCAmelCase      ,			self					).__init__()
     snake_case_         =   model.transformer.wte
     snake_case_         =   copy.deepcopy(model.transformer.wte.weight					)
     snake_case_         =   nn.Convad(self.embeddings.weight.size(1					)      ,			2_5_6      ,			3      ,			padding=1					)
     snake_case_         =   nn.Sequential(nn.Linear(2_5_6      ,			3_2					)      ,			nn.Dropout(p=0.1					)      ,			nn.Linear(3_2      ,			3_2					)      ,			nn.Linear(3_2      ,			1					)					)



  def    lowerCAmelCase__					(       self    :    Union[str, Any]      ,			_lowerCAmelCase    :    Optional[int]					)				->	List[Any]:




     """simple docstring"""
     snake_case_         =   torch.max(self.conv(self.embeddings(_lowerCAmelCase					).squeeze(1					).transpose(1      ,			2					)					)      ,			2					)[0]
     snake_case_         =   self.fc(_lowerCAmelCase					)
     return qs.squeeze(1					)



  @classmethod
  def    lowerCAmelCase__					(       cls    :    List[Any]      ,			_lowerCAmelCase    :    List[Any]      ,			_lowerCAmelCase    :    Dict					)				->	Dict:




     """simple docstring"""
     snake_case_         =   cls(_lowerCAmelCase					)  # this calls __init__
     snake_case_         =   torch.load(_lowerCAmelCase					)
     secondary_learner.load_state_dict(_lowerCAmelCase					)
     snake_case_         =   model.transformer.wte
     snake_case_         =   copy.deepcopy(model.transformer.wte.weight					)
     return secondary_learner




