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

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


__SCREAMING_SNAKE_CASE								=logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE								={
    "facebook/bart-large": "https://huggingface.co/facebook/bart-large/resolve/main/config.json",
    # See all BART models at https://huggingface.co/models?filter=bart
}







class  UpperCamelCase	(    lowercase_						):
   lowercase							=						'bart'
   lowercase							=						['past_key_values']
   lowercase							=						{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}
   def __init__(    self  ,__UpperCamelCase=5_0265  ,__UpperCamelCase=1024  ,__UpperCamelCase=12  ,__UpperCamelCase=4096  ,__UpperCamelCase=16  ,__UpperCamelCase=12  ,__UpperCamelCase=4096  ,__UpperCamelCase=16  ,__UpperCamelCase=0.0  ,__UpperCamelCase=0.0  ,__UpperCamelCase="gelu"  ,__UpperCamelCase=1024  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.0  ,__UpperCamelCase=0.0  ,__UpperCamelCase=0.02  ,__UpperCamelCase=0.0  ,__UpperCamelCase=False  ,__UpperCamelCase=True  ,__UpperCamelCase=3  ,__UpperCamelCase=1  ,__UpperCamelCase=0  ,__UpperCamelCase=2  ,__UpperCamelCase=True  ,__UpperCamelCase=2  ,__UpperCamelCase=2  ,**__UpperCamelCase  ,)			->   Tuple:


        '''simple docstring'''


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

        super().__init__(
            num_labels=__UpperCamelCase  ,pad_token_id=__UpperCamelCase  ,bos_token_id=__UpperCamelCase  ,eos_token_id=__UpperCamelCase  ,is_encoder_decoder=__UpperCamelCase  ,decoder_start_token_id=__UpperCamelCase  ,forced_eos_token_id=__UpperCamelCase  ,**__UpperCamelCase  ,)

        # ensure backward compatibility for BART CNN models
        if self.forced_bos_token_id is None and kwargs.get('force_bos_token_to_be_generated'  ,__UpperCamelCase    ):
             lowercase_			:       Optional[Any]									=			self.bos_token_id
             warnings.warn(
                 f'''Please make sure the config includes `forced_bos_token_id={self.bos_token_id}` in future versions. '''
                 'The config can simply be saved and uploaded again to be fixed.'    )







class  UpperCamelCase	(    lowercase_						):
   @property
   def     _UpperCAmelCase     (    self    )			->   Mapping[str, Mapping[int, str]]:


        '''simple docstring'''


        if self.task in ["default", "seq2seq-lm"]:
             lowercase_			:       Union[str, Any]									=			OrderedDict(
                 [
                     ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                     ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                 ]    )

             if self.use_past:
                  lowercase_			:       Optional[Any]									=			{0: 'batch'}
                  lowercase_			:       Optional[int]									=			{0: 'batch', 1: 'past_decoder_sequence + sequence'}
             else:
                  lowercase_			:       Optional[int]									=			{0: 'batch', 1: 'decoder_sequence'}
                  lowercase_			:       Optional[Any]									=			{0: 'batch', 1: 'decoder_sequence'}

             if self.use_past:
                  self.fill_with_past_key_values_(__UpperCamelCase  ,direction='inputs'    )
        elif self.task == "causal-lm":
             # TODO: figure this case out.
             lowercase_			:       Union[str, Any]									=			OrderedDict(
                 [
                     ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
                     ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
                 ]    )
             if self.use_past:
                  lowercase_			:       Union[str, Any]									=			self.num_layers
                  for i in range(__UpperCamelCase    ):
                       lowercase_			:       Union[str, Any]									=			{0: 'batch', 2: 'past_sequence + sequence'}
                       lowercase_			:       List[Any]									=			{0: 'batch', 2: 'past_sequence + sequence'}
        else:
             lowercase_			:       List[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
   def     _UpperCAmelCase     (    self    )			->   Mapping[str, Mapping[int, str]]:


        '''simple docstring'''


        if self.task in ["default", "seq2seq-lm"]:
             lowercase_			:       Union[str, Any]									=			super().outputs
        else:
             lowercase_			:       Tuple									=			super(__UpperCamelCase  ,self    ).outputs
             if self.use_past:
                  lowercase_			:       Optional[Any]									=			self.num_layers
                  for i in range(__UpperCamelCase    ):
                       lowercase_			:       Any									=			{0: 'batch', 2: 'past_sequence + sequence'}
                       lowercase_			:       List[str]									=			{0: 'batch', 2: 'past_sequence + sequence'}
        return common_outputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = -1  ,__UpperCamelCase = -1  ,__UpperCamelCase = False  ,__UpperCamelCase = None  ,)			->   Mapping[str, Any]:


        '''simple docstring'''


        lowercase_			:       List[str]									=			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            __UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )

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

        if self.use_past:
             if not is_torch_available():
                  raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'    )
             else:
                  import torch
             lowercase_			:       Optional[int]									=			common_inputs['input_ids'].shape
             lowercase_			:       str									=			common_inputs['decoder_input_ids'].shape[1]
             lowercase_			:       List[str]									=			self.num_attention_heads
             lowercase_			:       Dict									=			(
                 batch,
                 num_encoder_attention_heads,
                 encoder_seq_length,
                 self._config.hidden_size // num_encoder_attention_heads,
             )
             lowercase_			:       Any									=			decoder_seq_length + 3
             lowercase_			:       List[Any]									=			(
                 batch,
                 num_decoder_attention_heads,
                 decoder_past_length,
                 self._config.hidden_size // num_decoder_attention_heads,
             )

             lowercase_			:       List[str]									=			torch.cat(
                 [common_inputs['decoder_attention_mask'], torch.ones(__UpperCamelCase  ,__UpperCamelCase    )]  ,dim=1    )

             lowercase_			:       Optional[int]									=			[]
             # If the number of encoder and decoder layers are present in the model configuration, both are considered
             lowercase_			:       str									=			self.num_layers
             lowercase_			:       Optional[Any]									=			min(__UpperCamelCase  ,__UpperCamelCase    )
             lowercase_			:       Any									=			max(__UpperCamelCase  ,__UpperCamelCase    ) - min_num_layers
             lowercase_			:       Union[str, Any]									=			'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

             for _ in range(__UpperCamelCase    ):
                  common_inputs["past_key_values"].append(
                      (
                          torch.zeros(__UpperCamelCase    ),
                          torch.zeros(__UpperCamelCase    ),
                          torch.zeros(__UpperCamelCase    ),
                          torch.zeros(__UpperCamelCase    ),
                      )    )
             # TODO: test this.
             lowercase_			:       List[Any]									=			encoder_shape if remaining_side_name == 'encoder' else decoder_shape
             for _ in range(__UpperCamelCase  ,__UpperCamelCase    ):
                  common_inputs["past_key_values"].append((torch.zeros(__UpperCamelCase    ), torch.zeros(__UpperCamelCase    ))    )
        return common_inputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = -1  ,__UpperCamelCase = -1  ,__UpperCamelCase = False  ,__UpperCamelCase = None  ,)			->   Mapping[str, Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
            __UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )

        if self.use_past:
             if not is_torch_available():
                  raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'    )
             else:
                  import torch
             lowercase_			:       Dict									=			common_inputs['input_ids'].shape
             # Not using the same length for past_key_values
             lowercase_			:       Union[str, Any]									=			seqlen + 2
             lowercase_			:       Tuple									=			self.num_layers
             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_			:       List[Any]									=			common_inputs['attention_mask'].dtype
             lowercase_			:       Optional[Any]									=			torch.cat(
                 [common_inputs['attention_mask'], torch.ones(__UpperCamelCase  ,__UpperCamelCase  ,dtype=__UpperCamelCase    )]  ,dim=1    )
             lowercase_			:       Any									=			[
                 (torch.zeros(__UpperCamelCase    ), torch.zeros(__UpperCamelCase    )) for _ in range(__UpperCamelCase    )
             ]
        return common_inputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = -1  ,__UpperCamelCase = -1  ,__UpperCamelCase = False  ,__UpperCamelCase = None  ,)			->   Mapping[str, Any]:


        '''simple docstring'''


        lowercase_			:       List[Any]									=			compute_effective_axis_dimension(
            __UpperCamelCase  ,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_			:       int									=			tokenizer.num_special_tokens_to_add(__UpperCamelCase    )
        lowercase_			:       Optional[Any]									=			compute_effective_axis_dimension(
            __UpperCamelCase  ,fixed_dimension=OnnxConfig.default_fixed_sequence  ,num_token_to_add=__UpperCamelCase    )

        # Generate dummy inputs according to compute batch and sequence
        lowercase_			:       Tuple									=			[' '.join([tokenizer.unk_token]    ) * seq_length] * batch_size
        lowercase_			:       int									=			dict(tokenizer(__UpperCamelCase  ,return_tensors=__UpperCamelCase    )    )
        return common_inputs
   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase = -1  ,__UpperCamelCase = -1  ,__UpperCamelCase = False  ,__UpperCamelCase = None  ,)			->   Mapping[str, Any]:


        '''simple docstring'''


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

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

        return common_inputs







   def     _UpperCAmelCase     (    self  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase  ,__UpperCamelCase    )			->   int:


        '''simple docstring'''


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


"""simple docstring"""







import unittest

import numpy as np

from transformers import DistilBertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
       import jax.numpy as jnp

       from transformers.models.distilbert.modeling_flax_distilbert import (
           FlaxDistilBertForMaskedLM,
           FlaxDistilBertForMultipleChoice,
           FlaxDistilBertForQuestionAnswering,
           FlaxDistilBertForSequenceClassification,
           FlaxDistilBertForTokenClassification,
           FlaxDistilBertModel,
       )







class  UpperCamelCase	(    unittest.TestCase						):
   def __init__(    self  ,__UpperCamelCase  ,__UpperCamelCase=13  ,__UpperCamelCase=7  ,__UpperCamelCase=True  ,__UpperCamelCase=True  ,__UpperCamelCase=True  ,__UpperCamelCase=True  ,__UpperCamelCase=99  ,__UpperCamelCase=32  ,__UpperCamelCase=5  ,__UpperCamelCase=4  ,__UpperCamelCase=37  ,__UpperCamelCase="gelu"  ,__UpperCamelCase=0.1  ,__UpperCamelCase=0.1  ,__UpperCamelCase=512  ,__UpperCamelCase=16  ,__UpperCamelCase=2  ,__UpperCamelCase=0.02  ,__UpperCamelCase=4  ,)			->   str:


        '''simple docstring'''


        lowercase_			:       str									=			parent
        lowercase_			:       List[str]									=			batch_size
        lowercase_			:       Optional[Any]									=			seq_length
        lowercase_			:       str									=			is_training
        lowercase_			:       Any									=			use_attention_mask
        lowercase_			:       int									=			use_token_type_ids
        lowercase_			:       Union[str, Any]									=			use_labels
        lowercase_			:       int									=			vocab_size
        lowercase_			:       Union[str, Any]									=			hidden_size
        lowercase_			:       Optional[int]									=			num_hidden_layers
        lowercase_			:       Dict									=			num_attention_heads
        lowercase_			:       Dict									=			intermediate_size
        lowercase_			:       str									=			hidden_act
        lowercase_			:       Union[str, Any]									=			hidden_dropout_prob
        lowercase_			:       Dict									=			attention_probs_dropout_prob
        lowercase_			:       Union[str, Any]									=			max_position_embeddings
        lowercase_			:       List[Any]									=			type_vocab_size
        lowercase_			:       str									=			type_sequence_label_size
        lowercase_			:       Optional[Any]									=			initializer_range
        lowercase_			:       Any									=			num_choices
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size    )

        lowercase_			:       List[Any]									=			None
        if self.use_attention_mask:
             lowercase_			:       Dict									=			random_attention_mask([self.batch_size, self.seq_length]    )

        lowercase_			:       Optional[Any]									=			DistilBertConfig(
            vocab_size=self.vocab_size  ,dim=self.hidden_size  ,n_layers=self.num_hidden_layers  ,n_heads=self.num_attention_heads  ,hidden_dim=self.intermediate_size  ,hidden_act=self.hidden_act  ,dropout=self.hidden_dropout_prob  ,attention_dropout=self.attention_probs_dropout_prob  ,max_position_embeddings=self.max_position_embeddings  ,initializer_range=self.initializer_range  ,tie_weights_=__UpperCamelCase  ,)

        return config, input_ids, attention_mask







   def     _UpperCAmelCase     (    self    )			->   Dict:


        '''simple docstring'''


        lowercase_			:       Any									=			self.prepare_config_and_inputs()
        lowercase_							,						lowercase_							,						lowercase_			:       str									=			config_and_inputs
        lowercase_			:       List[str]									=			{'input_ids': input_ids, 'attention_mask': attention_mask}
        return config, inputs_dict







@require_flax
class  UpperCamelCase	(    lowercase_		,   unittest.TestCase						):
   lowercase							=						(
       (
           FlaxDistilBertModel,
           FlaxDistilBertForMaskedLM,
           FlaxDistilBertForMultipleChoice,
           FlaxDistilBertForQuestionAnswering,
           FlaxDistilBertForSequenceClassification,
           FlaxDistilBertForTokenClassification,
           FlaxDistilBertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   def     _UpperCAmelCase     (    self    )			->   List[Any]:


        '''simple docstring'''


        lowercase_			:       Optional[Any]									=			FlaxDistilBertModelTester(self    )







   @slow
   def     _UpperCAmelCase     (    self    )			->   List[str]:


        '''simple docstring'''


        for model_class_name in self.all_model_classes:
             lowercase_			:       Optional[int]									=			model_class_name.from_pretrained('distilbert-base-uncased'    )
             lowercase_			:       List[str]									=			model(np.ones((1, 1)    )    )
             self.assertIsNotNone(__UpperCamelCase    )







@require_flax
class  UpperCamelCase	(    unittest.TestCase						):
   @slow
   def     _UpperCAmelCase     (    self    )			->   str:


        '''simple docstring'''


        lowercase_			:       Optional[int]									=			FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'    )
        lowercase_			:       Tuple									=			np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]    )
        lowercase_			:       Optional[int]									=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]    )
        lowercase_			:       Any									=			model(__UpperCamelCase  ,attention_mask=__UpperCamelCase    )[0]
        lowercase_			:       Tuple									=			(1, 11, 768)
        self.assertEqual(output.shape  ,__UpperCamelCase    )
        lowercase_			:       Union[str, Any]									=			np.array([[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]]    )

        self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]  ,__UpperCamelCase  ,atol=1e-4    )    )

