






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



				"""simple docstring"""



				def __init__(      self    :    List[str]			,		UpperCamelCase__    :    Optional[Any]			,		UpperCamelCase__    :    Any=1_3			,		UpperCamelCase__    :    Tuple=7			,		UpperCamelCase__    :    Optional[Any]=True			,		UpperCamelCase__    :    int=True			,		UpperCamelCase__    :    Optional[Any]=True			,		UpperCamelCase__    :    List[str]=True			,		UpperCamelCase__    :    Any=9_9			,		UpperCamelCase__    :    Union[str, Any]=3_2			,		UpperCamelCase__    :    Dict=5			,		UpperCamelCase__    :    Tuple=4			,		UpperCamelCase__    :    Dict=3_7			,		UpperCamelCase__    :    Any="gelu"			,		UpperCamelCase__    :    Optional[int]=0.1			,		UpperCamelCase__    :    List[str]=0.1			,		UpperCamelCase__    :    int=5_1_2			,		UpperCamelCase__    :    int=1_6			,		UpperCamelCase__    :    str=2			,		UpperCamelCase__    :    List[str]=0.02			,		UpperCamelCase__    :    Optional[Any]=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 					__magic_name__							(      self    :    str):






											'''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_=lowerCAmelCase_			,		)

											return config, input_ids, attention_mask



				def 					__magic_name__							(      self    :    Union[str, Any]):






											'''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    _lowerCAmelCase			(		__lowerCAmelCase     ,      unittest.TestCase      ):



				"""simple docstring"""



				_lowercase				:   Dict      =  (
				    (
				        FlaxDistilBertModel,
				        FlaxDistilBertForMaskedLM,
				        FlaxDistilBertForMultipleChoice,
				        FlaxDistilBertForQuestionAnswering,
				        FlaxDistilBertForSequenceClassification,
				        FlaxDistilBertForTokenClassification,
				        FlaxDistilBertForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)



				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											snake_case__	  =       FlaxDistilBertModelTester(self)



				@slow
				def 					__magic_name__							(      self    :    str):






											'''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(lowerCAmelCase_)


@require_flax
class    _lowerCAmelCase			(		unittest.TestCase      ):



				"""simple docstring"""



				@slow
				def 					__magic_name__							(      self    :    Dict):






											'''simple docstring'''







											snake_case__	  =       FlaxDistilBertModel.from_pretrained("""distilbert-base-uncased""")
											snake_case__	  =       np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]])
											snake_case__	  =       np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
											snake_case__	  =       model(lowerCAmelCase_			,		attention_mask=lowerCAmelCase_)[0]
											snake_case__	  =       (1, 1_1, 7_6_8)
											self.assertEqual(output.shape			,		lowerCAmelCase_)
											snake_case__	  =       np.array([[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]])

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








from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_torch_available, is_torch_tpu_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_torch_available():
	import torch

if is_torch_tpu_available(check_device=False):
	import torch_xla.core.xla_model as xm


a__       =     logging.get_logger(__name__)


@dataclass
class    _lowerCAmelCase			(		lowercase_      ):



				"""simple docstring"""



				_lowercase				:   Dict      =  [
				    '''no_inference''',
				    '''no_cuda''',
				    '''no_tpu''',
				    '''no_speed''',
				    '''no_memory''',
				    '''no_env_print''',
				    '''no_multi_process''',
				]



				def __init__(      self    :    str			,		**UpperCamelCase__    :    Dict):






											'''simple docstring'''







											for deprecated_arg in self.deprecated_args:
																		if deprecated_arg in kwargs:
																									snake_case__	  =       deprecated_arg[3:]
																									setattr(self			,		UpperCamelCase__			,		not kwargs.pop(UpperCamelCase__))
																									logger.warning(
																									    F'''{deprecated_arg} is depreciated. Please use --no_{positive_arg} or'''
																									    F''' {positive_arg}={kwargs[positive_arg]}''')

											snake_case__	  =       kwargs.pop("""torchscript"""			,		self.torchscript)
											snake_case__	  =       kwargs.pop("""torch_xla_tpu_print_metrics"""			,		self.torch_xla_tpu_print_metrics)
											snake_case__	  =       kwargs.pop("""fp16_opt_level"""			,		self.fpaa_opt_level)
											super().__init__(**UpperCamelCase__)

				_lowercase				:   bool      =  field(default=lowercase_     ,      metadata={'''help''': '''Trace the models using torchscript'''}      )
				_lowercase				:   bool      =  field(default=lowercase_     ,      metadata={'''help''': '''Print Xla/PyTorch tpu metrics'''}      )
				_lowercase				:   str      =  field(
				    default='''O1'''     ,      metadata={
				        '''help''': (
				            '''For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\']. '''
				            '''See details at https://nvidia.github.io/apex/amp.html'''
				        )
				    }     ,      )



				@cached_property
				def 					__magic_name__							(      self    :    Tuple):






											'''simple docstring'''







											requires_backends(self			,		["""torch"""])
											logger.info("""PyTorch: setting up devices""")
											if not self.cuda:
																		snake_case__	  =       torch.device("""cpu""")
																		snake_case__	  =       0
											elif is_torch_tpu_available():
																		snake_case__	  =       xm.xla_device()
																		snake_case__	  =       0
											else:
																		snake_case__	  =       torch.device("""cuda""" if torch.cuda.is_available() else """cpu""")
																		snake_case__	  =       torch.cuda.device_count()
											return device, n_gpu



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






											'''simple docstring'''







											return is_torch_tpu_available() and self.tpu



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






											'''simple docstring'''







											requires_backends(self			,		["""torch"""])
											# TODO(PVP): currently only single GPU is supported
											return torch.cuda.current_device()



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






											'''simple docstring'''







											requires_backends(self			,		["""torch"""])
											return self._setup_devices[0]



				@property
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											requires_backends(self			,		["""torch"""])
											return self._setup_devices[1]



				@property
				def 					__magic_name__							(      self    :    str):






											'''simple docstring'''







											return self.n_gpu > 0
