






'''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


__lowerCAmelCase					   =       logging.get_logger(__name__)

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




class      __SCREAMING_SNAKE_CASE					(__A					):


  """simple docstring"""


  _a     :   List[Any]			 =    RetriBertConfig
  _a     :   Tuple			 =    None
  _a     :   Dict			 =    '''retribert'''







  def       _a					( self      ,  UpperCamelCase__  ):






         """simple docstring"""


         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  )





__lowerCAmelCase					   =       R"\n\n    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the\n    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads\n    etc.)\n\n    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.\n    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage\n    and behavior.\n\n    Parameters:\n        config ([`RetriBertConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"




@add_start_docstrings(
    '''Bert Based model to embed queries or document for document retrieval.'''       ,					__A       ,					)
class      __SCREAMING_SNAKE_CASE					(__A					):


  """simple docstring"""


  def __init__( self      ,  UpperCamelCase__  ):






         """simple docstring"""


         super().__init__(UpperCamelCase__  )
         a_	    = config.projection_dim

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

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

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







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__=-1      ,  ):






         """simple docstring"""


         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
                a_	    = input_ids.device
                a_	    = input_ids.size()
                a_	    = torch.zeros(UpperCamelCase__      ,  dtype=torch.long      ,  device=UpperCamelCase__  )
                a_	    = [None] * sent_encoder.config.num_hidden_layers
                a_	    = sent_encoder.get_extended_attention_mask(
                    UpperCamelCase__      ,  UpperCamelCase__  )

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

                # run embedding layer on everything at once
                a_	    = 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
                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(UpperCamelCase__      ,  UpperCamelCase__      ,  UpperCamelCase__  )
                       pooled_output_list.append(UpperCamelCase__  )
                return torch.cat(UpperCamelCase__      ,  dim=0  )







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__=None      ,  UpperCamelCase__=-1      ,  ):






         """simple docstring"""


         a_	    = self.embed_sentences_checkpointed(
             UpperCamelCase__      ,  UpperCamelCase__      ,  self.bert_query      ,  UpperCamelCase__      ,  )
         return self.project_query(UpperCamelCase__  )







  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__=None      ,  UpperCamelCase__=-1      ,  ):






         """simple docstring"""


         a_	    = 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      ,  ):






         """simple docstring"""


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












'''simple docstring'''







import itertools
import os
import subprocess
from os.path import dirname

from parameterized import parameterized

from tests.trainer.test_trainer import TrainerIntegrationCommon  # noqa
from transformers import is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_tests_dir,
    require_deepspeed,
    require_torch_gpu,
    slow,
)
from transformers.trainer_utils import set_seed


if is_torch_available():
  from tests.trainer.test_trainer import (  # noqa
      RegressionModelConfig,
      RegressionPreTrainedModel,
      get_regression_trainer,
  )


set_seed(42)

__lowerCAmelCase					   =       get_tests_dir("fixtures")
__lowerCAmelCase					   =       os.path.join(dirname(get_tests_dir()))
__lowerCAmelCase					   =       dirname(os.path.abspath(__file__))

# default torch.distributed port
__lowerCAmelCase					   =       "10999"

__lowerCAmelCase					   =       "t5-small"

# *** Working Models ***
__lowerCAmelCase					   =       "hf-internal-testing/tiny-albert"
__lowerCAmelCase					   =       "sshleifer/bart-tiny-random"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-bert"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-bigbird_pegasus"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-big_bird"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-blenderbot"
__lowerCAmelCase					   =       "bigscience/bigscience-small-testing"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-deberta"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-deberta-v2"
__lowerCAmelCase					   =       "sshleifer/tiny-distilbert-base-cased"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-electra"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-flaubert"
__lowerCAmelCase					   =       "stas/tiny-wmt19-en-de"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-funnel"
__lowerCAmelCase					   =       "sshleifer/tiny-gpt2"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-gptj"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-gpt_neo"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-layoutlm"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-led"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-longformer"
__lowerCAmelCase					   =       "stas/tiny-m2m_100"  # hf tiny model is unsuitable
__lowerCAmelCase					   =       "sshleifer/tiny-marian-en-de"
__lowerCAmelCase					   =       "sshleifer/tiny-mbart"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-mobilebert"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-mpnet"
__lowerCAmelCase					   =       "stas/pegasus-cnn_dailymail-tiny-random"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-prophetnet"
__lowerCAmelCase					   =       "sshleifer/tiny-distilroberta-base"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-squeezebert"
__lowerCAmelCase					   =       "patrickvonplaten/t5-tiny-random"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-t5-v1.1"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-vit"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-xlm-roberta"
__lowerCAmelCase					   =       "sshleifer/tiny-xlnet-base-cased"


# *** To Fix ***


# *** tiny model issues ***
# missing model files:
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-mt5"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-camembert"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-openai-gpt"

# missing tokenizer files
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-convbert"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-layoutlmv2"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-hubert"

# issues with tokenizer
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-ctrl"
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-transfo-xl"  # same as ctrl

# other issues with tiny models
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-ibert"  # multiple issues with either mlm/qa/clas
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-reformer"  # multiple issues with either mlm/qa/clas

# *** Lacking official examples to test with ***
# or not working with examples
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-dpr"
# - "dpr"  examples/research_projects/rag-end2end-retriever/
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-rag"
# - "rag" research_projects
__lowerCAmelCase					   =       ""
# - "luke" Entities classes - no plan to make such example
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-lxmert"
# - "lxmert" doesn't work with run_qa.py
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-clip"
# - "clip" nothing under pytorch examples - XXX: Suraj is working on adding some - check by end of Sep
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-speech_to_text"
# - "speech_to_text", nothing under pytorch examples


# *** Reactive mode ***
# models with low usage, unstable API, things about to change - do nothing about the following until someone runs into a problem
__lowerCAmelCase					   =       "hf-internal-testing/tiny-random-tapas"


def 			__UpperCamelCase				(			lowercase_					:     Optional[int]=False      ):
       """simple docstring"""




       a_	    = min(2      ,					get_gpu_count()      ) if distributed else 1
       a_	    = os.environ.get('DS_TEST_PORT'      ,					lowercase_      )
       return F'deepspeed --num_nodes 1 --num_gpus {num_gpus} --master_port {master_port}'.split()


def 			__UpperCamelCase				(			):
       """simple docstring"""




       a_	    = F'{FIXTURE_DIRECTORY}/tests_samples'
       a_	    = F'{data_dir_samples}/wmt_en_ro'
       a_	    = F'{data_dir_samples}/xsum'
       a_	    = '\n        --do_train\n        --max_train_samples 4\n        --per_device_train_batch_size 2\n        --num_train_epochs 1\n        --fp16\n        --report_to none\n        --overwrite_output_dir\n        '.split()

       # try to cover as many models as possible once (it's enough to run on one task per model)
       # but need a tiny model for each
       #
       # should have "{model_type.upper()}_TINY" corresponding vars defined, e.g., T5_TINY, etc.
       a_	    = {
           'trans': [
               'bart',
               'fsmt',
               'm2m_100',
               'marian',
               'mbart',
               't5',
               't5_v1',
               # "mt5", missing model files
           ],
           'sum': [
               'pegasus',
           ],
           'clm': [
               'big_bird',
               'bigbird_pegasus',
               'blenderbot',
               'bloom',
               'gpt2',
               'gpt_neo',
               'gptj',
               'xlm-roberta',
               'prophetnet',
               # "camembert", missing model files
           ],
           'mlm': [
               'albert',
               'deberta',
               'deberta-v2',
               'distilbert',
               'electra',
               'flaubert',
               'funnel',
               'layoutlm',
               # "reformer", # multiple issues with either mlm/qa/clas
           ],
           'qa': [
               'led',
               'longformer',
               'mobilebert',
               'mpnet',
               'roberta',
               'squeezebert',
               # "convbert", # missing tokenizer files
               # "layoutlmv2", missing model files
           ],
           'clas': [
               'bert',
               'xlnet',
               # "hubert", # missing tokenizer files
               # "ibert", # multiple issues with either mlm/qa/clas
               # "transfo-xl", # tokenizer issues as ctrl
               # "ctrl", # tokenizer issues
               # "openai-gpt", missing model files
               # "tapas", multiple issues
           ],
           'img_clas': [
               'vit',
           ],
       }

       a_	    = F'{ROOT_DIRECTORY}/examples/pytorch'

       a_	    = {
           'trans': F'\n        {scripts_dir}/translation/run_translation.py\n        --train_file {data_dir_wmt}/train.json\n        --source_lang en\n        --target_lang ro\n        ',
           'sum': F'\n        {scripts_dir}/summarization/run_summarization.py\n        --train_file {data_dir_xsum}/sample.json\n        --max_source_length 12\n        --max_target_length 12\n        --lang en\n        ',
           'clm': F'\n        {scripts_dir}/language-modeling/run_clm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        --block_size 8\n        ',
           'mlm': F'\n        {scripts_dir}/language-modeling/run_mlm.py\n        --train_file {FIXTURE_DIRECTORY}/sample_text.txt\n        ',
           'qa': F'\n        {scripts_dir}/question-answering/run_qa.py\n        --train_file {data_dir_samples}/SQUAD/sample.json\n        ',
           'clas': F'\n        {scripts_dir}/text-classification/run_glue.py\n        --train_file {data_dir_samples}/MRPC/train.csv\n        --max_seq_length 12\n        --task_name MRPC\n        ',
           'img_clas': F'\n        {scripts_dir}/image-classification/run_image_classification.py\n            --dataset_name hf-internal-testing/cats_vs_dogs_sample\n            --remove_unused_columns False\n            --max_steps 10\n            --image_processor_name {DS_TESTS_DIRECTORY}/vit_feature_extractor.json\n        ',
       }

       a_	    = get_launcher(distributed=lowercase_      )

       a_	    = {}
       for task, args in tasks.items():
              a_	    = args.split()
              for model in tasksamodels[task]:
                     a_	    = globals()[F'{model.upper().replace("-"      ,					"_"      )}_TINY']
                     a_	    = F'--model_name_or_path {model_name}'.split()
                     a_	    = launcher + args + args_model + args_main

                     # # generation special case
                     # if task == "gen":
                     #     launcher = f"deepspeed --num_nodes 1 --num_gpus 1".split()
                     #     args_model += f"--model_type {model}".split()
                     #     cmds[f"{task}_{model}"] = launcher + args + args_model
                     # else:

       return cmds


__lowerCAmelCase					   =       make_task_cmds()

__lowerCAmelCase					   =       "zero2"
__lowerCAmelCase					   =       "zero3"

__lowerCAmelCase					   =       [ZEROa, ZEROa]


def 			__UpperCamelCase				(			lowercase_					:     str      ,					lowercase_					:     Optional[int]      ,					lowercase_					:     Dict      ):
       """simple docstring"""




       a_	    = parameterized.to_safe_name('_'.join(str(lowercase_      ) for x in param.args      )      )
       return F'{func.__name__}_{param_based_name}'


# Cartesian-product of zero stages with models to test
__lowerCAmelCase					   =       list(itertools.product(stages, task_cmds.keys()))




@slow
@require_deepspeed
@require_torch_gpu
class      __SCREAMING_SNAKE_CASE					(__A					):


  """simple docstring"""


  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         if task not in task_cmds:
                raise ValueError(f'don\'t know of task {task}, have {task_cmds.keys()}'  )

         a_	    = task_cmds[task]
         a_	    = f'--deepspeed {self.test_file_dir_str}/ds_config_{stage}.json'.split()

         a_	    = self.get_auto_remove_tmp_dir()
         a_	    = f'--output_dir {output_dir}'.split()

         cmd += args_ds + args_out

         return cmd, output_dir







  @parameterized.expand(UpperCamelCase__      ,  name_func=UpperCamelCase__  )
  def       _a					( self      ,  UpperCamelCase__      ,  UpperCamelCase__  ):






         """simple docstring"""


         a_    ,  a_	    = self.get_task_cmd(UpperCamelCase__      ,  UpperCamelCase__  )

         # 1. generate the checkpoint
         cmd += "--save_steps 1".split()
         # keep for quick debug
         # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] + cmd)); die
         execute_subprocess_async(UpperCamelCase__      ,  env=self.get_env()  )

         # 2. test that the fp32 weights get reconsolidated
         a_	    = f'{output_dir}/checkpoint-1'
         a_	    = f'{chkpt_dir}/out.bin'
         a_	    = f'{chkpt_dir}/zero_to_fp32.py {chkpt_dir} {recovered_model_path}'
         # keep for quick debug
         # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
         subprocess.check_call(UpperCamelCase__      ,  shell=UpperCamelCase__  )
         assert os.path.exists(UpperCamelCase__  ), f'{recovered_model_path} was not found'

         # possibly could also test that the resulting saved model is usable but given that we use
         # random models we won't know if it's any good




