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





	'''simple docstring'''



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


				'''simple docstring'''


				_snake_case		=			parent
				_snake_case		=			batch_size
				_snake_case		=			seq_length
				_snake_case		=			is_training
				_snake_case		=			use_attention_mask
				_snake_case		=			use_token_type_ids
				_snake_case		=			use_labels
				_snake_case		=			vocab_size
				_snake_case		=			hidden_size
				_snake_case		=			num_hidden_layers
				_snake_case		=			num_attention_heads
				_snake_case		=			intermediate_size
				_snake_case		=			hidden_act
				_snake_case		=			hidden_dropout_prob
				_snake_case		=			attention_probs_dropout_prob
				_snake_case		=			max_position_embeddings
				_snake_case		=			type_vocab_size
				_snake_case		=			type_sequence_label_size
				_snake_case		=			initializer_range
				_snake_case		=			num_choices



	def  A							(  self	:    Optional[int]	):


				'''simple docstring'''


				_snake_case		=			ids_tensor([self.batch_size, self.seq_length]			,  self.vocab_size	)

				_snake_case		=			None
				if self.use_attention_mask:
							_snake_case		=			random_attention_mask([self.batch_size, self.seq_length]	)

				_snake_case		=			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_=lowercase			,  )

				return config, input_ids, attention_mask





	def  A							(  self	:    List[str]	):


				'''simple docstring'''


				_snake_case		=			self.prepare_config_and_inputs()
				_snake_case	,					_snake_case	,					_snake_case		=			config_and_inputs
				_snake_case		=			{'input_ids': input_ids, 'attention_mask': attention_mask}
				return config, inputs_dict




@require_flax
class    SCREAMING_SNAKE_CASE__     (	UpperCAmelCase							,unittest.TestCase					):





	'''simple docstring'''



	_UpperCAmelCase       :		List[Any]       =					(
	    (
	        FlaxDistilBertModel,
	        FlaxDistilBertForMaskedLM,
	        FlaxDistilBertForMultipleChoice,
	        FlaxDistilBertForQuestionAnswering,
	        FlaxDistilBertForSequenceClassification,
	        FlaxDistilBertForTokenClassification,
	        FlaxDistilBertForQuestionAnswering,
	    )
	    if is_flax_available()
	    else ()
	)



	def  A							(  self	:    Tuple	):


				'''simple docstring'''


				_snake_case		=			FlaxDistilBertModelTester(self	)





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


				'''simple docstring'''


				for model_class_name in self.all_model_classes:
							_snake_case		=			model_class_name.from_pretrained('distilbert-base-uncased'	)
							_snake_case		=			model(np.ones((1, 1)	)	)
							self.assertIsNotNone(lowercase	)




@require_flax
class    SCREAMING_SNAKE_CASE__     (	unittest.TestCase					):





	'''simple docstring'''



	@slow
	def  A							(  self	:    str	):


				'''simple docstring'''


				_snake_case		=			FlaxDistilBertModel.from_pretrained('distilbert-base-uncased'	)
				_snake_case		=			np.array([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]]	)
				_snake_case		=			np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]	)
				_snake_case		=			model(lowercase			,  attention_mask=lowercase	)[0]
				_snake_case		=			(1, 11, 768)
				self.assertEqual(output.shape			,  lowercase	)
				_snake_case		=			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]			,  lowercase			,  atol=1E-4	)	)

from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import TensorType, logging


if TYPE_CHECKING:
			from ...onnx.config import PatchingSpec
			from ...tokenization_utils_base import PreTrainedTokenizerBase


_lowerCamelCase       :				List[Any]       =			logging.get_logger(__name__)

_lowerCamelCase       :				Union[str, Any]       =			{
    '''allenai/longformer-base-4096''': '''https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096''': '''https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json''',
    '''allenai/longformer-large-4096-finetuned-triviaqa''': (
        '''https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json'''
    ),
    '''allenai/longformer-base-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
    '''allenai/longformer-large-4096-extra.pos.embd.only''': (
        '''https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json'''
    ),
}




class    SCREAMING_SNAKE_CASE__     (	UpperCAmelCase					):





	'''simple docstring'''



	_UpperCAmelCase       :		Dict       =					"longformer"



	def __init__(  self	:    Optional[Any]			,  lowercase	:    Union[List[int], int] = 512			,  lowercase	:    int = 2			,  lowercase	:    int = 1			,  lowercase	:    int = 0			,  lowercase	:    int = 2			,  lowercase	:    int = 30_522			,  lowercase	:    int = 768			,  lowercase	:    int = 12			,  lowercase	:    int = 12			,  lowercase	:    int = 3_072			,  lowercase	:    str = "gelu"			,  lowercase	:    float = 0.1			,  lowercase	:    float = 0.1			,  lowercase	:    int = 512			,  lowercase	:    int = 2			,  lowercase	:    float = 0.02			,  lowercase	:    float = 1E-12			,  lowercase	:    bool = False			,  **lowercase	:    Optional[Any]			,  ):


				'''simple docstring'''


				super().__init__(pad_token_id=lowercase			,  **lowercase	)

				_snake_case		=			attention_window
				_snake_case		=			sep_token_id
				_snake_case		=			bos_token_id
				_snake_case		=			eos_token_id
				_snake_case		=			vocab_size
				_snake_case		=			hidden_size
				_snake_case		=			num_hidden_layers
				_snake_case		=			num_attention_heads
				_snake_case		=			hidden_act
				_snake_case		=			intermediate_size
				_snake_case		=			hidden_dropout_prob
				_snake_case		=			attention_probs_dropout_prob
				_snake_case		=			max_position_embeddings
				_snake_case		=			type_vocab_size
				_snake_case		=			initializer_range
				_snake_case		=			layer_norm_eps
				_snake_case		=			onnx_export




class    SCREAMING_SNAKE_CASE__     (	UpperCAmelCase					):





	'''simple docstring'''



	def __init__(  self	:    int			,  lowercase	:    "PretrainedConfig"			,  lowercase	:    str = "default"			,  lowercase	:    "List[PatchingSpec]" = None	):


				'''simple docstring'''


				super().__init__(lowercase			,  lowercase			,  lowercase	)
				_snake_case		=			True



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


				'''simple docstring'''


				if self.task == "multiple-choice":
							_snake_case		=			{0: 'batch', 1: 'choice', 2: 'sequence'}
				else:
							_snake_case		=			{0: 'batch', 1: 'sequence'}
				return OrderedDict(
				    [
				        ('input_ids', dynamic_axis),
				        ('attention_mask', dynamic_axis),
				        ('global_attention_mask', dynamic_axis),
				    ]	)



	@property
	def  A							(  self	:    int	):


				'''simple docstring'''


				_snake_case		=			super().outputs
				if self.task == "default":
							_snake_case		=			{0: 'batch'}
				return outputs



	@property
	def  A							(  self	:    List[Any]	):


				'''simple docstring'''


				return 1E-4



	@property
	def  A							(  self	:    List[str]	):


				'''simple docstring'''


				return max(super().default_onnx_opset			,  14	)





	def  A							(  self	:    str			,  lowercase	:    "PreTrainedTokenizerBase"			,  lowercase	:    int = -1			,  lowercase	:    int = -1			,  lowercase	:    bool = False			,  lowercase	:    Optional[TensorType] = None			,  ):


				'''simple docstring'''


				_snake_case		=			super().generate_dummy_inputs(
				    preprocessor=lowercase			,  batch_size=lowercase			,  seq_length=lowercase			,  is_pair=lowercase			,  framework=lowercase	)
				import torch

				# for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
				# makes the export fail randomly
				_snake_case		=			torch.zeros_like(inputs['input_ids']	)
				# make every second token global
				_snake_case		=			1

				return inputs