




'''simple docstring'''


import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import (
     RemBertConfig,
     RemBertForCausalLM,
     RemBertForMaskedLM,
     RemBertForMultipleChoice,
     RemBertForQuestionAnswering,
     RemBertForSequenceClassification,
     RemBertForTokenClassification,
     RemBertModel,
 )
 from transformers.models.rembert.modeling_rembert import REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST
class    a_     :





     def __init__(				self	:      List[str]			,   lowercase	:      Tuple			,   lowercase	:      Union[str, Any]=13			,   lowercase	:      List[Any]=7			,   lowercase	:      Tuple=True			,   lowercase	:      str=True			,   lowercase	:      Optional[int]=True			,   lowercase	:      Tuple=True			,   lowercase	:      Any=99			,   lowercase	:      Optional[Any]=32			,   lowercase	:      str=18			,   lowercase	:      int=43			,   lowercase	:      Optional[int]=5			,   lowercase	:      List[str]=4			,   lowercase	:      str=37			,   lowercase	:      Optional[int]="gelu"			,   lowercase	:      str=0.1			,   lowercase	:      List[Any]=0.1			,   lowercase	:      Union[str, Any]=512			,   lowercase	:      Dict=16			,   lowercase	:      Optional[int]=2			,   lowercase	:      Dict=0.02			,   lowercase	:      Any=3			,   lowercase	:      Tuple=4			,   lowercase	:      Any=None			,   ):





        """simple docstring"""




        lowercase_    :Dict					    =				parent
        lowercase_    :Optional[int]					    =				batch_size
        lowercase_    :int					    =				seq_length
        lowercase_    :str					    =				is_training
        lowercase_    :Dict					    =				use_input_mask
        lowercase_    :List[Any]					    =				use_token_type_ids
        lowercase_    :List[Any]					    =				use_labels
        lowercase_    :str					    =				vocab_size
        lowercase_    :Tuple					    =				hidden_size
        lowercase_    :Union[str, Any]					    =				input_embedding_size
        lowercase_    :Any					    =				output_embedding_size
        lowercase_    :List[Any]					    =				num_hidden_layers
        lowercase_    :Optional[Any]					    =				num_attention_heads
        lowercase_    :Tuple					    =				intermediate_size
        lowercase_    :Any					    =				hidden_act
        lowercase_    :Optional[int]					    =				hidden_dropout_prob
        lowercase_    :Tuple					    =				attention_probs_dropout_prob
        lowercase_    :List[str]					    =				max_position_embeddings
        lowercase_    :str					    =				type_vocab_size
        lowercase_    :Tuple					    =				type_sequence_label_size
        lowercase_    :int					    =				initializer_range
        lowercase_    :List[str]					    =				num_labels
        lowercase_    :Any					    =				num_choices
        lowercase_    :Union[str, Any]					    =				scope





     def 							lowercase__   (				self	:      int		):





        """simple docstring"""




        lowercase_    :List[str]					    =				ids_tensor([self.batch_size, self.seq_length]			,   self.vocab_size		)

        lowercase_    :Union[str, Any]					    =				None
        if self.use_input_mask:
           lowercase_    :List[str]					    =				random_attention_mask([self.batch_size, self.seq_length]		)

        lowercase_    :Union[str, Any]					    =				None
        if self.use_token_type_ids:
           lowercase_    :Optional[int]					    =				ids_tensor([self.batch_size, self.seq_length]			,   self.type_vocab_size		)

        lowercase_    :Dict					    =				None
        lowercase_    :Any					    =				None
        lowercase_    :int					    =				None
        if self.use_labels:
           lowercase_    :Dict					    =				ids_tensor([self.batch_size]			,   self.type_sequence_label_size		)
           lowercase_    :List[str]					    =				ids_tensor([self.batch_size, self.seq_length]			,   self.num_labels		)
           lowercase_    :Dict					    =				ids_tensor([self.batch_size]			,   self.num_choices		)

        lowercase_    :List[Any]					    =				RemBertConfig(
            vocab_size=self.vocab_size			,   hidden_size=self.hidden_size			,   input_embedding_size=self.input_embedding_size			,   output_embedding_size=self.output_embedding_size			,   num_hidden_layers=self.num_hidden_layers			,   num_attention_heads=self.num_attention_heads			,   intermediate_size=self.intermediate_size			,   hidden_act=self.hidden_act			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   max_position_embeddings=self.max_position_embeddings			,   type_vocab_size=self.type_vocab_size			,   is_decoder=_lowerCAmelCase			,   initializer_range=self.initializer_range			,   )

        return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





     def 							lowercase__   (				self	:      List[Any]		):





        """simple docstring"""




        (
            (
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							    
        )    :str					    =				self.prepare_config_and_inputs()

        lowercase_    :Tuple					    =				True
        lowercase_    :Tuple					    =				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
        lowercase_    :int					    =				ids_tensor([self.batch_size, self.seq_length]			,   vocab_size=2		)

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )





     def 							lowercase__   (				self	:      int			,   lowercase	:      str			,   lowercase	:      Optional[int]			,   lowercase	:      Optional[int]			,   lowercase	:      Optional[int]			,   lowercase	:      Optional[int]			,   lowercase	:      List[str]			,   lowercase	:      Any		):





        """simple docstring"""




        lowercase_    :List[Any]					    =				RemBertModel(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :int					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase		)
        lowercase_    :Tuple					    =				model(_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase		)
        lowercase_    :int					    =				model(_lowerCAmelCase		)
        self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, self.seq_length, self.hidden_size)		)





     def 							lowercase__   (				self	:      List[str]			,   lowercase	:      Optional[int]			,   lowercase	:      Union[str, Any]			,   lowercase	:      Any			,   lowercase	:      Optional[Any]			,   lowercase	:      Optional[Any]			,   lowercase	:      int			,   lowercase	:      List[str]			,   lowercase	:      Dict			,   lowercase	:      int			,   ):





        """simple docstring"""




        lowercase_    :Optional[int]					    =				True
        lowercase_    :Any					    =				RemBertModel(_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :str					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   encoder_hidden_states=_lowerCAmelCase			,   encoder_attention_mask=_lowerCAmelCase			,   )
        lowercase_    :Optional[int]					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   encoder_hidden_states=_lowerCAmelCase			,   )
        lowercase_    :Union[str, Any]					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase		)
        self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, self.seq_length, self.hidden_size)		)





     def 							lowercase__   (				self	:      List[Any]			,   lowercase	:      List[str]			,   lowercase	:      Union[str, Any]			,   lowercase	:      Dict			,   lowercase	:      Union[str, Any]			,   lowercase	:      Tuple			,   lowercase	:      Optional[int]			,   lowercase	:      Optional[Any]			,   lowercase	:      str			,   lowercase	:      Optional[int]			,   ):





        """simple docstring"""




        lowercase_    :Optional[Any]					    =				RemBertForCausalLM(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :Optional[Any]					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   labels=_lowerCAmelCase		)
        self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)		)





     def 							lowercase__   (				self	:      str			,   lowercase	:      str			,   lowercase	:      Any			,   lowercase	:      Optional[Any]			,   lowercase	:      Any			,   lowercase	:      str			,   lowercase	:      Optional[int]			,   lowercase	:      int		):





        """simple docstring"""




        lowercase_    :Optional[int]					    =				RemBertForMaskedLM(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :Tuple					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   labels=_lowerCAmelCase		)
        self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.vocab_size)		)





     def 							lowercase__   (				self	:      List[Any]			,   lowercase	:      List[Any]			,   lowercase	:      str			,   lowercase	:      Dict			,   lowercase	:      Any			,   lowercase	:      Tuple			,   lowercase	:      int			,   lowercase	:      Any			,   lowercase	:      Optional[Any]			,   lowercase	:      List[str]			,   ):





        """simple docstring"""




        lowercase_    :Tuple					    =				True
        lowercase_    :str					    =				True
        lowercase_    :Optional[int]					    =				RemBertForCausalLM(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()

        # first forward pass
        lowercase_    :List[str]					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   encoder_hidden_states=_lowerCAmelCase			,   encoder_attention_mask=_lowerCAmelCase			,   use_cache=_lowerCAmelCase			,   )
        lowercase_    :str					    =				outputs.past_key_values

        # create hypothetical multiple next token and extent to next_input_ids
        lowercase_    :Union[str, Any]					    =				ids_tensor((self.batch_size, 3)			,   config.vocab_size		)
        lowercase_    :Optional[int]					    =				ids_tensor((self.batch_size, 3)			,   vocab_size=2		)

        # append to next input_ids and
        lowercase_    :List[Any]					    =				torch.cat([input_ids, next_tokens]			,   dim=-1		)
        lowercase_    :Optional[Any]					    =				torch.cat([input_mask, next_mask]			,   dim=-1		)

        lowercase_    :Optional[int]					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   encoder_hidden_states=_lowerCAmelCase			,   encoder_attention_mask=_lowerCAmelCase			,   output_hidden_states=_lowerCAmelCase			,   )["hidden_states"][0]
        lowercase_    :str					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   encoder_hidden_states=_lowerCAmelCase			,   encoder_attention_mask=_lowerCAmelCase			,   past_key_values=_lowerCAmelCase			,   output_hidden_states=_lowerCAmelCase			,   )["hidden_states"][0]

        # select random slice
        lowercase_    :Optional[Any]					    =				ids_tensor((1,)			,   output_from_past.shape[-1]		).item()
        lowercase_    :Optional[int]					    =				output_from_no_past[:, -3:, random_slice_idx].detach()
        lowercase_    :str					    =				output_from_past[:, :, random_slice_idx].detach()

        self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(_lowerCAmelCase			,   _lowerCAmelCase			,   atol=1e-3		)		)





     def 							lowercase__   (				self	:      str			,   lowercase	:      Tuple			,   lowercase	:      int			,   lowercase	:      Optional[int]			,   lowercase	:      Any			,   lowercase	:      List[str]			,   lowercase	:      Tuple			,   lowercase	:      Union[str, Any]		):





        """simple docstring"""




        lowercase_    :str					    =				RemBertForQuestionAnswering(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :List[str]					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   start_positions=_lowerCAmelCase			,   end_positions=_lowerCAmelCase			,   )
        self.parent.assertEqual(result.start_logits.shape			,   (self.batch_size, self.seq_length)		)
        self.parent.assertEqual(result.end_logits.shape			,   (self.batch_size, self.seq_length)		)





     def 							lowercase__   (				self	:      Optional[int]			,   lowercase	:      Any			,   lowercase	:      Any			,   lowercase	:      List[str]			,   lowercase	:      int			,   lowercase	:      str			,   lowercase	:      List[str]			,   lowercase	:      Dict		):





        """simple docstring"""




        lowercase_    :Tuple					    =				self.num_labels
        lowercase_    :Union[str, Any]					    =				RemBertForSequenceClassification(_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :List[str]					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   labels=_lowerCAmelCase		)
        self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_labels)		)





     def 							lowercase__   (				self	:      int			,   lowercase	:      Optional[int]			,   lowercase	:      str			,   lowercase	:      str			,   lowercase	:      int			,   lowercase	:      int			,   lowercase	:      Dict			,   lowercase	:      str		):





        """simple docstring"""




        lowercase_    :Optional[Any]					    =				self.num_labels
        lowercase_    :List[Any]					    =				RemBertForTokenClassification(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :int					    =				model(_lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   labels=_lowerCAmelCase		)
        self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.seq_length, self.num_labels)		)





     def 							lowercase__   (				self	:      str			,   lowercase	:      str			,   lowercase	:      Optional[Any]			,   lowercase	:      List[Any]			,   lowercase	:      List[Any]			,   lowercase	:      int			,   lowercase	:      int			,   lowercase	:      str		):





        """simple docstring"""




        lowercase_    :Union[str, Any]					    =				self.num_choices
        lowercase_    :Tuple					    =				RemBertForMultipleChoice(config=_lowerCAmelCase		)
        model.to(_lowerCAmelCase		)
        model.eval()
        lowercase_    :Optional[int]					    =				input_ids.unsqueeze(1		).expand(-1			,   self.num_choices			,   -1		).contiguous()
        lowercase_    :Optional[Any]					    =				token_type_ids.unsqueeze(1		).expand(-1			,   self.num_choices			,   -1		).contiguous()
        lowercase_    :Optional[int]					    =				input_mask.unsqueeze(1		).expand(-1			,   self.num_choices			,   -1		).contiguous()
        lowercase_    :Tuple					    =				model(
            _lowerCAmelCase			,   attention_mask=_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   labels=_lowerCAmelCase			,   )
        self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.num_choices)		)







     def 							lowercase__   (				self	:      Optional[Any]		):





        """simple docstring"""




        lowercase_    :str					    =				self.prepare_config_and_inputs()
        (
            (
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							    
        )    :Optional[Any]					    =				config_and_inputs
        lowercase_    :Dict					    =				{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_torch
class    a_     (       A__			,		A__			,		unittest.TestCase							):
     __A  				=    (
         (
             RemBertModel,
             RemBertForMaskedLM,
             RemBertForCausalLM,
             RemBertForMultipleChoice,
             RemBertForQuestionAnswering,
             RemBertForSequenceClassification,
             RemBertForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     __A  				=    (RemBertForCausalLM,) if is_torch_available() else ()
     __A  				=    (
         {
             "feature-extraction": RemBertModel,
             "fill-mask": RemBertForMaskedLM,
             "question-answering": RemBertForQuestionAnswering,
             "text-classification": RemBertForSequenceClassification,
             "text-generation": RemBertForCausalLM,
             "token-classification": RemBertForTokenClassification,
             "zero-shot": RemBertForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )





     def 							lowercase__   (				self	:      Tuple		):





        """simple docstring"""




        lowercase_    :List[Any]					    =				RemBertModelTester(self		)
        lowercase_    :Optional[Any]					    =				ConfigTester(self			,   config_class=_lowerCAmelCase			,   hidden_size=37		)





     def 							lowercase__   (				self	:      Tuple		):





        """simple docstring"""




        self.config_tester.run_common_tests()





     def 							lowercase__   (				self	:      Dict		):





        """simple docstring"""




        lowercase_    :Any					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      List[Any]		):





        """simple docstring"""




        lowercase_    :str					    =				self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
           lowercase_    :List[Any]					    =				type
           self.model_tester.create_and_check_model(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      int		):





        """simple docstring"""




        lowercase_    :Tuple					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      List[str]		):





        """simple docstring"""




        lowercase_    :Optional[Any]					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      int		):





        """simple docstring"""




        lowercase_    :Any					    =				self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      Tuple		):





        """simple docstring"""




        lowercase_    :str					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      Any		):





        """simple docstring"""




        lowercase_    :Tuple					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      Tuple		):





        """simple docstring"""




        lowercase_    :Tuple					    =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      int		):





        """simple docstring"""




        lowercase_    :Any					    =				self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*_lowerCAmelCase		)





     def 							lowercase__   (				self	:      str		):





        """simple docstring"""




        (
            (
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							(
            lowercase_    
        )    ,							    
        )    :Optional[Any]					    =				self.model_tester.prepare_config_and_inputs_for_decoder()

        lowercase_    :Dict					    =				None

        self.model_tester.create_and_check_model_as_decoder(
            _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   )







     @slow
     def 							lowercase__   (				self	:      Optional[int]		):





        """simple docstring"""




        for model_name in REMBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           lowercase_    :Tuple					    =				RemBertModel.from_pretrained(_lowerCAmelCase		)
           self.assertIsNotNone(_lowerCAmelCase		)







@require_torch
class    a_     (       unittest.TestCase							):





     @slow
     def 							lowercase__   (				self	:      List[str]		):





        """simple docstring"""




        lowercase_    :Tuple					    =				RemBertModel.from_pretrained("google/rembert"		)
        lowercase_    :Optional[Any]					    =				torch.tensor([[312, 56_498, 313, 2_125, 313]]		)
        lowercase_    :Optional[int]					    =				torch.tensor([[0, 0, 0, 1, 1]]		)
        with torch.no_grad():
           lowercase_    :Optional[Any]					    =				model(_lowerCAmelCase			,   token_type_ids=_lowerCAmelCase			,   output_hidden_states=_lowerCAmelCase		)

        lowercase_    :Optional[Any]					    =				1_152

        lowercase_    :int					    =				torch.Size((1, 5, hidden_size)		)
        self.assertEqual(output["last_hidden_state"].shape			,   _lowerCAmelCase		)

        lowercase_    :List[str]					    =				torch.tensor(
            [
                [
                    [0.07_54, -0.20_22, 0.19_04],
                    [-0.33_54, -0.36_92, -0.47_91],
                    [-0.23_14, -0.67_29, -0.07_49],
                    [-0.03_96, -0.31_05, -0.42_34],
                    [-0.15_71, -0.05_25, 0.53_53],
                ]
            ]		)

        # Running on the original tf implementation gives slightly different results here.
        # Not clear why this variations is present
        # TODO: Find reason for discrepancy
        # expected_original_implementation = [[
        #     [0.07630594074726105, -0.20146065950393677, 0.19107051193714142],
        #     [-0.3405614495277405, -0.36971670389175415, -0.4808273911476135],
        #     [-0.22587086260318756, -0.6656315922737122, -0.07844287157058716],
        #     [-0.04145475849509239, -0.3077218234539032, -0.42316967248916626],
        #     [-0.15887849032878876, -0.054529931396245956, 0.5356100797653198]
        # ]]

        self.assertTrue(torch.allclose(output["last_hidden_state"][:, :, :3]			,   _lowerCAmelCase			,   atol=1e-4		)		)









'''simple docstring'''


from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...file_utils import TensorType, is_torch_available
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import logging


lowerCAmelCase      :  str      							=logging.get_logger(__name__)

lowerCAmelCase      :  Union[str, Any]      							={
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}
class    a_     (       _lowerCAmelCase							):
     __A  				=    "blenderbot"
     __A  				=    ["past_key_values"]
     __A  				=    {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}





     def __init__(				self	:      Dict			,   lowercase	:      List[Any]=8_008			,   lowercase	:      List[str]=128			,   lowercase	:      Optional[int]=2			,   lowercase	:      Optional[int]=10_240			,   lowercase	:      Tuple=32			,   lowercase	:      int=24			,   lowercase	:      Dict=10_240			,   lowercase	:      Any=32			,   lowercase	:      Optional[int]=0.0			,   lowercase	:      Dict=0.0			,   lowercase	:      int=True			,   lowercase	:      Tuple=True			,   lowercase	:      List[str]="gelu"			,   lowercase	:      int=2_560			,   lowercase	:      Optional[Any]=0.1			,   lowercase	:      Any=0.0			,   lowercase	:      List[str]=0.0			,   lowercase	:      Optional[Any]=0.02			,   lowercase	:      Tuple=1			,   lowercase	:      Union[str, Any]=False			,   lowercase	:      str=0			,   lowercase	:      List[str]=1			,   lowercase	:      Optional[Any]=2			,   lowercase	:      Union[str, Any]=3			,   lowercase	:      List[str]=2			,   **lowercase	:      Optional[Any]			,   ):





        """simple docstring"""




        lowercase_    :str					    =				vocab_size
        lowercase_    :Any					    =				max_position_embeddings
        lowercase_    :List[Any]					    =				d_model
        lowercase_    :Union[str, Any]					    =				encoder_ffn_dim
        lowercase_    :str					    =				encoder_layers
        lowercase_    :Optional[int]					    =				encoder_attention_heads
        lowercase_    :Dict					    =				decoder_ffn_dim
        lowercase_    :str					    =				decoder_layers
        lowercase_    :int					    =				decoder_attention_heads
        lowercase_    :Union[str, Any]					    =				dropout
        lowercase_    :int					    =				attention_dropout
        lowercase_    :int					    =				activation_dropout
        lowercase_    :str					    =				activation_function
        lowercase_    :Tuple					    =				init_std
        lowercase_    :Any					    =				encoder_layerdrop
        lowercase_    :Dict					    =				decoder_layerdrop
        lowercase_    :int					    =				use_cache
        lowercase_    :Dict					    =				encoder_layers
        lowercase_    :Dict					    =				scale_embedding  # scale factor will be sqrt(d_model) if True

        super().__init__(
            pad_token_id=lowercase			,   bos_token_id=lowercase			,   eos_token_id=lowercase			,   is_encoder_decoder=lowercase			,   decoder_start_token_id=lowercase			,   encoder_no_repeat_ngram_size=lowercase			,   forced_eos_token_id=lowercase			,   **lowercase			,   )







class    a_     (       _lowerCAmelCase							):





     @property
     def 							lowercase__   (				self	:      Union[str, Any]		):





        """simple docstring"""




        if self.task in ["default", "seq2seq-lm"]:
           lowercase_    :Optional[int]					    =				OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
               ]		)
           if self.use_past:
              lowercase_    :int					    =				{0: "batch"}
              lowercase_    :List[str]					    =				{0: "batch", 1: "past_decoder_sequence + sequence"}
           else:
              lowercase_    :Union[str, Any]					    =				{0: "batch", 1: "decoder_sequence"}
              lowercase_    :List[Any]					    =				{0: "batch", 1: "decoder_sequence"}
           if self.use_past:
              self.fill_with_past_key_values_(lowercase			,   direction="inputs"		)
        elif self.task == "causal-lm":
           lowercase_    :List[Any]					    =				OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
               ]		)
           if self.use_past:
              lowercase_    ,							lowercase_    :Optional[int]					    =				self.num_layers
              for i in range(lowercase		):
                 lowercase_    :List[str]					    =				{0: "batch", 2: "past_sequence + sequence"}
                 lowercase_    :Any					    =				{0: "batch", 2: "past_sequence + sequence"}
        else:
           lowercase_    :str					    =				OrderedDict(
               [
                   ("input_ids", {0: "batch", 1: "encoder_sequence"}),
                   ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
                   ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
                   ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
               ]		)

        return common_inputs





     @property
     # Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
     def 							lowercase__   (				self	:      Dict		):





        """simple docstring"""




        if self.task in ["default", "seq2seq-lm"]:
           lowercase_    :Union[str, Any]					    =				super().outputs
        else:
           lowercase_    :Optional[Any]					    =				super(lowercase			,   self		).outputs
           if self.use_past:
              lowercase_    ,							lowercase_    :Any					    =				self.num_layers
              for i in range(lowercase		):
                 lowercase_    :str					    =				{0: "batch", 2: "past_sequence + sequence"}
                 lowercase_    :Any					    =				{0: "batch", 2: "past_sequence + sequence"}
        return common_outputs





     def 							lowercase__   (				self	:      Tuple			,   lowercase	:      PreTrainedTokenizer			,   lowercase	:      int = -1			,   lowercase	:      int = -1			,   lowercase	:      bool = False			,   lowercase	:      Optional[TensorType] = None			,   ):





        """simple docstring"""




        lowercase_    :int					    =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase			,   lowercase			,   lowercase			,   lowercase			,   lowercase		)
        # Generate decoder inputs
        lowercase_    :List[Any]					    =				seq_length if not self.use_past else 1
        lowercase_    :Optional[Any]					    =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase			,   lowercase			,   lowercase			,   lowercase			,   lowercase		)
        lowercase_    :Tuple					    =				{F'decoder_{name}': tensor for name, tensor in decoder_inputs.items()}
        lowercase_    :Any					    =				dict(**lowercase			,   **lowercase		)

        if self.use_past:
           if not is_torch_available():
              raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."		)
           else:
              import torch
           lowercase_    ,							lowercase_    :int					    =				common_inputs["input_ids"].shape
           lowercase_    :Tuple					    =				common_inputs["decoder_input_ids"].shape[1]
           lowercase_    ,							lowercase_    :Tuple					    =				self.num_attention_heads
           lowercase_    :List[Any]					    =				(
               batch,
               num_encoder_attention_heads,
               encoder_seq_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )
           lowercase_    :List[Any]					    =				decoder_seq_length
           lowercase_    :str					    =				(
               batch,
               num_decoder_attention_heads,
               decoder_past_length,
               self._config.hidden_size // num_decoder_attention_heads,
           )
           lowercase_    :Optional[int]					    =				torch.cat(
               [common_inputs["decoder_attention_mask"], torch.ones(lowercase			,   lowercase		)]			,   dim=1		)
           lowercase_    :int					    =				[]
           lowercase_    ,							lowercase_    :Tuple					    =				self.num_layers

           for _ in range(lowercase		):
              common_inputs["past_key_values"].append(
                  (
                      torch.zeros(lowercase		),
                      torch.zeros(lowercase		),
                      torch.zeros(lowercase		),
                      torch.zeros(lowercase		),
                  )		)
        return common_inputs





     def 							lowercase__   (				self	:      Optional[Any]			,   lowercase	:      PreTrainedTokenizer			,   lowercase	:      int = -1			,   lowercase	:      int = -1			,   lowercase	:      bool = False			,   lowercase	:      Optional[TensorType] = None			,   ):





        """simple docstring"""




        lowercase_    :str					    =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            lowercase			,   lowercase			,   lowercase			,   lowercase			,   lowercase		)

        if self.use_past:
           if not is_torch_available():
              raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."		)
           else:
              import torch
           lowercase_    ,							lowercase_    :List[str]					    =				common_inputs["input_ids"].shape
           lowercase_    :int					    =				seqlen
           lowercase_    ,							lowercase_    :int					    =				self.num_layers
           lowercase_    ,							lowercase_    :Union[str, Any]					    =				self.num_attention_heads
           lowercase_    :Any					    =				(
               batch,
               num_encoder_attention_heads,
               past_key_values_length,
               self._config.hidden_size // num_encoder_attention_heads,
           )
           lowercase_    :Optional[Any]					    =				common_inputs["attention_mask"].dtype
           lowercase_    :List[Any]					    =				torch.cat(
               [common_inputs["attention_mask"], torch.ones(lowercase			,   lowercase			,   dtype=lowercase		)]			,   dim=1		)
           lowercase_    :Any					    =				[
               (torch.zeros(lowercase		), torch.zeros(lowercase		)) for _ in range(lowercase		)
           ]
        return common_inputs





     def 							lowercase__   (				self	:      Dict			,   lowercase	:      PreTrainedTokenizer			,   lowercase	:      int = -1			,   lowercase	:      int = -1			,   lowercase	:      bool = False			,   lowercase	:      Optional[TensorType] = None			,   ):





        """simple docstring"""




        lowercase_    :Optional[int]					    =				compute_effective_axis_dimension(
            lowercase			,   fixed_dimension=OnnxConfig.default_fixed_batch			,   num_token_to_add=0		)

        # If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
        lowercase_    :Optional[Any]					    =				tokenizer.num_special_tokens_to_add(lowercase		)
        lowercase_    :int					    =				compute_effective_axis_dimension(
            lowercase			,   fixed_dimension=OnnxConfig.default_fixed_sequence			,   num_token_to_add=lowercase		)

        # Generate dummy inputs according to compute batch and sequence
        lowercase_    :int					    =				[" ".join([tokenizer.unk_token]		) * seq_length] * batch_size
        lowercase_    :Dict					    =				dict(tokenizer(lowercase			,   return_tensors=lowercase		)		)
        return common_inputs





     def 							lowercase__   (				self	:      Optional[Any]			,   lowercase	:      PreTrainedTokenizer			,   lowercase	:      int = -1			,   lowercase	:      int = -1			,   lowercase	:      bool = False			,   lowercase	:      Optional[TensorType] = None			,   ):





        """simple docstring"""




        if self.task in ["default", "seq2seq-lm"]:
           lowercase_    :str					    =				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
               lowercase			,   batch_size=lowercase			,   seq_length=lowercase			,   is_pair=lowercase			,   framework=lowercase		)

        elif self.task == "causal-lm":
           lowercase_    :Optional[Any]					    =				self._generate_dummy_inputs_for_causal_lm(
               lowercase			,   batch_size=lowercase			,   seq_length=lowercase			,   is_pair=lowercase			,   framework=lowercase		)
        else:
           lowercase_    :List[str]					    =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
               lowercase			,   batch_size=lowercase			,   seq_length=lowercase			,   is_pair=lowercase			,   framework=lowercase		)

        return common_inputs





     def 							lowercase__   (				self	:      Optional[int]			,   lowercase	:      str			,   lowercase	:      Union[str, Any]			,   lowercase	:      int			,   lowercase	:      List[Any]		):





        """simple docstring"""




        if self.task in ["default", "seq2seq-lm"]:
           lowercase_    :Optional[int]					    =				super()._flatten_past_key_values_(lowercase			,   lowercase			,   lowercase			,   lowercase		)
        else:
           lowercase_    :Tuple					    =				super(lowercase			,   self		)._flatten_past_key_values_(
               lowercase			,   lowercase			,   lowercase			,   lowercase		)







     def 							lowercase__   (				self	:      List[Any]			,   lowercase	:      Mapping[str, Mapping[int, str]]			,   lowercase	:      str		):





        """simple docstring"""




        if direction not in ["inputs", "outputs"]:
           raise ValueError(F'direction must either be "inputs" or "outputs", but {direction} was given'		)

        lowercase_    :Any					    =				"past_key_values" if direction == "inputs" else "present"
        lowercase_    ,							lowercase_    :Optional[int]					    =				self.num_layers

        lowercase_    :str					    =				"past_encoder_sequence"
        lowercase_    :Tuple					    =				"past_decoder_sequence" if direction == "inputs" else "past_decoder_sequence + sequence"

        for i in range(lowercase		):
           lowercase_    :Dict					    =				{0: "batch", 2: decoder_sequence}
           lowercase_    :Optional[int]					    =				{0: "batch", 2: decoder_sequence}
           lowercase_    :Optional[int]					    =				{0: "batch", 2: encoder_sequence}
           lowercase_    :int					    =				{0: "batch", 2: encoder_sequence}


