




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


UpperCamelCase__						=			logging.get_logger(__name__)

UpperCamelCase__						=			{
    '''facebook/blenderbot_small-90M''': '''https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/config.json''',
    # See all BlenderbotSmall models at https://huggingface.co/models?filter=blenderbot_small
}




class   a__			(       UpperCamelCase_  ):
						snake_case__             = '''blenderbot-small'''
						snake_case__             = ['''past_key_values''']
						snake_case__             = {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}





						def __init__(							self      :	Tuple     ,a__      :	List[Any]=5_0265     ,a__      :	Tuple=512     ,a__      :	Optional[int]=8     ,a__      :	Union[str, Any]=2048     ,a__      :	str=16     ,a__      :	Optional[Any]=8     ,a__      :	Any=2048     ,a__      :	List[str]=16     ,a__      :	Dict=0.0     ,a__      :	List[Any]=0.0     ,a__      :	str=True     ,a__      :	Optional[Any]=True     ,a__      :	Dict="gelu"     ,a__      :	int=512     ,a__      :	Union[str, Any]=0.1     ,a__      :	Union[str, Any]=0.0     ,a__      :	Union[str, Any]=0.0     ,a__      :	Union[str, Any]=0.02     ,a__      :	int=1     ,a__      :	Any=False     ,a__      :	List[str]=0     ,a__      :	Optional[Any]=1     ,a__      :	List[str]=2     ,a__      :	Dict=2     ,**a__      :	List[Any]     ,)						->       str:




											"""simple docstring"""







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

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




class   a__			(       UpperCamelCase_  ):





						@property
						def       __UpperCamelCase     (							self      :	str)						->       Mapping[str, Mapping[int, str]]:




											"""simple docstring"""







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

																if self.use_past:
																					_lowerCAmelCase:List[str]									=  {0: '''batch'''}
																					_lowerCAmelCase:Any									=  {0: '''batch''', 1: '''past_decoder_sequence + sequence'''}
																else:
																					_lowerCAmelCase:str									=  {0: '''batch''', 1: '''decoder_sequence'''}
																					_lowerCAmelCase:List[str]									=  {0: '''batch''', 1: '''decoder_sequence'''}

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




											"""simple docstring"""







											if self.task in ["default", "seq2seq-lm"]:
																_lowerCAmelCase:Tuple									=  super().outputs
											else:
																_lowerCAmelCase:List[Any]									=  super(a__     ,self).outputs
																if self.use_past:
																					_lowerCAmelCase							,   _lowerCAmelCase:int									=  self.num_layers
																					for i in range(a__):
																										_lowerCAmelCase:List[Any]									=  {0: '''batch''', 2: '''past_sequence + sequence'''}
																										_lowerCAmelCase:int									=  {0: '''batch''', 2: '''past_sequence + sequence'''}
											return common_outputs





						def       __UpperCamelCase     (							self      :	Dict     ,a__      :	PreTrainedTokenizer     ,a__      :	int = -1     ,a__      :	int = -1     ,a__      :	bool = False     ,a__      :	Optional[TensorType] = None     ,)						->       Mapping[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:Optional[int]									=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
											    a__     ,a__     ,a__     ,a__     ,a__)

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

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

																_lowerCAmelCase:Any									=  torch.cat(
																    [common_inputs['''decoder_attention_mask'''], torch.ones(a__     ,a__)]     ,dim=1)

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

																for _ in range(a__):
																					common_inputs["past_key_values"].append(
																					    (
																					        torch.zeros(a__),
																					        torch.zeros(a__),
																					        torch.zeros(a__),
																					        torch.zeros(a__),
																					    ))
																# TODO: test this.
																_lowerCAmelCase:Optional[Any]									=  encoder_shape if remaining_side_name == '''encoder''' else decoder_shape
																for _ in range(a__     ,a__):
																					common_inputs["past_key_values"].append((torch.zeros(a__), torch.zeros(a__)))
											return common_inputs





						def       __UpperCamelCase     (							self      :	int     ,a__      :	PreTrainedTokenizer     ,a__      :	int = -1     ,a__      :	int = -1     ,a__      :	bool = False     ,a__      :	Optional[TensorType] = None     ,)						->       Mapping[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:List[Any]									=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
											    a__     ,a__     ,a__     ,a__     ,a__)

											if self.use_past:
																if not is_torch_available():
																					raise ValueError('''Cannot generate dummy past_keys inputs without PyTorch installed.''')
																else:
																					import torch
																_lowerCAmelCase							,   _lowerCAmelCase:List[str]									=  common_inputs['''input_ids'''].shape
																# Not using the same length for past_key_values
																_lowerCAmelCase:str									=  seqlen + 2
																_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  self.num_layers
																_lowerCAmelCase							,   _lowerCAmelCase:str									=  self.num_attention_heads
																_lowerCAmelCase:List[Any]									=  (
																    batch,
																    num_encoder_attention_heads,
																    past_key_values_length,
																    self._config.hidden_size // num_encoder_attention_heads,
																)

																_lowerCAmelCase:List[Any]									=  common_inputs['''attention_mask'''].dtype
																_lowerCAmelCase:int									=  torch.cat(
																    [common_inputs['''attention_mask'''], torch.ones(a__     ,a__     ,dtype=a__)]     ,dim=1)
																_lowerCAmelCase:List[str]									=  [
																    (torch.zeros(a__), torch.zeros(a__)) for _ in range(a__)
																]
											return common_inputs





						def       __UpperCamelCase     (							self      :	int     ,a__      :	PreTrainedTokenizer     ,a__      :	int = -1     ,a__      :	int = -1     ,a__      :	bool = False     ,a__      :	Optional[TensorType] = None     ,)						->       Mapping[str, Any]:




											"""simple docstring"""







											_lowerCAmelCase:Union[str, Any]									=  compute_effective_axis_dimension(
											    a__     ,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
											_lowerCAmelCase:Optional[Any]									=  tokenizer.num_special_tokens_to_add(a__)
											_lowerCAmelCase:Dict									=  compute_effective_axis_dimension(
											    a__     ,fixed_dimension=OnnxConfig.default_fixed_sequence     ,num_token_to_add=a__)

											# Generate dummy inputs according to compute batch and sequence
											_lowerCAmelCase:Union[str, Any]									=  [''' '''.join([tokenizer.unk_token]) * seq_length] * batch_size
											_lowerCAmelCase:int									=  dict(tokenizer(a__     ,return_tensors=a__))
											return common_inputs





						def       __UpperCamelCase     (							self      :	Optional[Any]     ,a__      :	PreTrainedTokenizer     ,a__      :	int = -1     ,a__      :	int = -1     ,a__      :	bool = False     ,a__      :	Optional[TensorType] = None     ,)						->       Mapping[str, Any]:




											"""simple docstring"""







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

											elif self.task == "causal-lm":
																_lowerCAmelCase:Dict									=  self._generate_dummy_inputs_for_causal_lm(
																    a__     ,batch_size=a__     ,seq_length=a__     ,is_pair=a__     ,framework=a__)
											else:
																_lowerCAmelCase:Tuple									=  self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
																    a__     ,batch_size=a__     ,seq_length=a__     ,is_pair=a__     ,framework=a__)

											return common_inputs







						def       __UpperCamelCase     (							self      :	List[str]     ,a__      :	List[Any]     ,a__      :	str     ,a__      :	Optional[int]     ,a__      :	Any)						->       List[Any]:




											"""simple docstring"""







											if self.task in ["default", "seq2seq-lm"]:
																_lowerCAmelCase:Dict									=  super()._flatten_past_key_values_(a__     ,a__     ,a__     ,a__)
											else:
																_lowerCAmelCase:List[str]									=  super(a__     ,self)._flatten_past_key_values_(
																    a__     ,a__     ,a__     ,a__)












"""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   a__			(       unittest.TestCase  ):





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




											"""simple docstring"""







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





						def       __UpperCamelCase     (							self      :	List[str])						->       int:




											"""simple docstring"""







											_lowerCAmelCase:int									=  ids_tensor([self.batch_size, self.seq_length]     ,self.vocab_size)

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

											_lowerCAmelCase:Dict									=  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_=a__     ,)

											return config, input_ids, attention_mask







						def       __UpperCamelCase     (							self      :	Any)						->       Any:




											"""simple docstring"""







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




@require_flax
class   a__			(       UpperCamelCase_					,	unittest.TestCase  ):
						snake_case__             = (
						    (
						        FlaxDistilBertModel,
						        FlaxDistilBertForMaskedLM,
						        FlaxDistilBertForMultipleChoice,
						        FlaxDistilBertForQuestionAnswering,
						        FlaxDistilBertForSequenceClassification,
						        FlaxDistilBertForTokenClassification,
						        FlaxDistilBertForQuestionAnswering,
						    )
						    if is_flax_available()
						    else ()
						)





						def       __UpperCamelCase     (							self      :	str)						->       Dict:




											"""simple docstring"""







											_lowerCAmelCase:str									=  FlaxDistilBertModelTester(self)







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




											"""simple docstring"""







											for model_class_name in self.all_model_classes:
																_lowerCAmelCase:List[str]									=  model_class_name.from_pretrained('''distilbert-base-uncased''')
																_lowerCAmelCase:Tuple									=  model(np.ones((1, 1)))
																self.assertIsNotNone(a__)




@require_flax
class   a__			(       unittest.TestCase  ):





						@slow
						def       __UpperCamelCase     (							self      :	Union[str, Any])						->       Optional[Any]:




											"""simple docstring"""







											_lowerCAmelCase:Any									=  FlaxDistilBertModel.from_pretrained('''distilbert-base-uncased''')
											_lowerCAmelCase:Optional[int]									=  np.array([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]])
											_lowerCAmelCase:Tuple									=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
											_lowerCAmelCase:List[str]									=  model(a__     ,attention_mask=a__)[0]
											_lowerCAmelCase:List[str]									=  (1, 11, 768)
											self.assertEqual(output.shape     ,a__)
											_lowerCAmelCase:int									=  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]     ,a__     ,atol=1E-4))






