





'''simple docstring'''






import os
import tempfile
import unittest

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

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


if is_torch_available():
      import torch

      from transformers import (
          DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          DistilBertForMaskedLM,
          DistilBertForMultipleChoice,
          DistilBertForQuestionAnswering,
          DistilBertForSequenceClassification,
          DistilBertForTokenClassification,
          DistilBertModel,
      )




class  UpperCAmelCase     (			snake_case_   ):




      def __init__(			self			: Union[str, Any]			,			__snake_case			: List[Any]			,			__snake_case			: Union[str, Any]=13			,			__snake_case			: List[str]=7			,			__snake_case			: List[Any]=True			,			__snake_case			: Union[str, Any]=True			,			__snake_case			: str=False			,			__snake_case			: str=True			,			__snake_case			: Union[str, Any]=99			,			__snake_case			: Union[str, Any]=32			,			__snake_case			: Optional[int]=5			,			__snake_case			: Dict=4			,			__snake_case			: Union[str, Any]=37			,			__snake_case			: Tuple="gelu"			,			__snake_case			: Tuple=0.1			,			__snake_case			: List[str]=0.1			,			__snake_case			: Optional[Any]=5_12			,			__snake_case			: Tuple=16			,			__snake_case			: Optional[int]=2			,			__snake_case			: Tuple=0.02			,			__snake_case			: Optional[int]=3			,			__snake_case			: str=4			,			__snake_case			: Any=None			,			)				->      Union[str, Any]:
            _lowerCAmelCase							     =							parent
            _lowerCAmelCase							     =							batch_size
            _lowerCAmelCase							     =							seq_length
            _lowerCAmelCase							     =							is_training
            _lowerCAmelCase							     =							use_input_mask
            _lowerCAmelCase							     =							use_token_type_ids
            _lowerCAmelCase							     =							use_labels
            _lowerCAmelCase							     =							vocab_size
            _lowerCAmelCase							     =							hidden_size
            _lowerCAmelCase							     =							num_hidden_layers
            _lowerCAmelCase							     =							num_attention_heads
            _lowerCAmelCase							     =							intermediate_size
            _lowerCAmelCase							     =							hidden_act
            _lowerCAmelCase							     =							hidden_dropout_prob
            _lowerCAmelCase							     =							attention_probs_dropout_prob
            _lowerCAmelCase							     =							max_position_embeddings
            _lowerCAmelCase							     =							type_vocab_size
            _lowerCAmelCase							     =							type_sequence_label_size
            _lowerCAmelCase							     =							initializer_range
            _lowerCAmelCase							     =							num_labels
            _lowerCAmelCase							     =							num_choices
            _lowerCAmelCase							     =							scope




      def 				lowercase__    (			self			: Tuple  )				->      Optional[Any]:
            _lowerCAmelCase							     =							ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size  )

            _lowerCAmelCase							     =							None
            if self.use_input_mask:
                  _lowerCAmelCase							     =							random_attention_mask([self.batch_size, self.seq_length]  )

            _lowerCAmelCase							     =							None
            _lowerCAmelCase							     =							None
            _lowerCAmelCase							     =							None
            if self.use_labels:
                  _lowerCAmelCase							     =							ids_tensor([self.batch_size]			,			self.type_sequence_label_size  )
                  _lowerCAmelCase							     =							ids_tensor([self.batch_size, self.seq_length]			,			self.num_labels  )
                  _lowerCAmelCase							     =							ids_tensor([self.batch_size]			,			self.num_choices  )

            _lowerCAmelCase							     =							self.get_config()

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




      def 				lowercase__    (			self			: List[str]  )				->      Tuple:
            return 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			,			)




      def 				lowercase__    (			self			: str			,			__snake_case			: str			,			__snake_case			: List[Any]			,			__snake_case			: Any			,			__snake_case			: Tuple			,			__snake_case			: Union[str, Any]			,			__snake_case			: str  )				->      Any:
            _lowerCAmelCase							     =							DistilBertModel(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()
            _lowerCAmelCase							     =							model(__snake_case			,			__snake_case  )
            _lowerCAmelCase							     =							model(__snake_case  )
            self.parent.assertEqual(result.last_hidden_state.shape			,			(self.batch_size, self.seq_length, self.hidden_size)  )




      def 				lowercase__    (			self			: int			,			__snake_case			: List[Any]			,			__snake_case			: Union[str, Any]			,			__snake_case			: Optional[int]			,			__snake_case			: Dict			,			__snake_case			: Tuple			,			__snake_case			: Tuple  )				->      Optional[int]:
            _lowerCAmelCase							     =							DistilBertForMaskedLM(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()
            _lowerCAmelCase							     =							model(__snake_case			,			attention_mask=__snake_case			,			labels=__snake_case  )
            self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.vocab_size)  )




      def 				lowercase__    (			self			: Union[str, Any]			,			__snake_case			: str			,			__snake_case			: Optional[int]			,			__snake_case			: Any			,			__snake_case			: str			,			__snake_case			: Any			,			__snake_case			: Any  )				->      int:
            _lowerCAmelCase							     =							DistilBertForQuestionAnswering(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()
            _lowerCAmelCase							     =							model(
                __snake_case			,			attention_mask=__snake_case			,			start_positions=__snake_case			,			end_positions=__snake_case  )
            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]			,			__snake_case			: str			,			__snake_case			: Union[str, Any]			,			__snake_case			: Optional[Any]			,			__snake_case			: Union[str, Any]			,			__snake_case			: Dict			,			__snake_case			: List[Any]  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							self.num_labels
            _lowerCAmelCase							     =							DistilBertForSequenceClassification(__snake_case  )
            model.to(__snake_case  )
            model.eval()
            _lowerCAmelCase							     =							model(__snake_case			,			attention_mask=__snake_case			,			labels=__snake_case  )
            self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_labels)  )




      def 				lowercase__    (			self			: Optional[Any]			,			__snake_case			: List[Any]			,			__snake_case			: Dict			,			__snake_case			: List[str]			,			__snake_case			: str			,			__snake_case			: Optional[Any]			,			__snake_case			: Tuple  )				->      str:
            _lowerCAmelCase							     =							self.num_labels
            _lowerCAmelCase							     =							DistilBertForTokenClassification(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()

            _lowerCAmelCase							     =							model(__snake_case			,			attention_mask=__snake_case			,			labels=__snake_case  )
            self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.seq_length, self.num_labels)  )




      def 				lowercase__    (			self			: Optional[Any]			,			__snake_case			: Any			,			__snake_case			: Optional[Any]			,			__snake_case			: Any			,			__snake_case			: List[str]			,			__snake_case			: Optional[int]			,			__snake_case			: Tuple  )				->      int:
            _lowerCAmelCase							     =							self.num_choices
            _lowerCAmelCase							     =							DistilBertForMultipleChoice(config=__snake_case  )
            model.to(__snake_case  )
            model.eval()
            _lowerCAmelCase							     =							input_ids.unsqueeze(1  ).expand(-1			,			self.num_choices			,			-1  ).contiguous()
            _lowerCAmelCase							     =							input_mask.unsqueeze(1  ).expand(-1			,			self.num_choices			,			-1  ).contiguous()
            _lowerCAmelCase							     =							model(
                __snake_case			,			attention_mask=__snake_case			,			labels=__snake_case			,			)
            self.parent.assertEqual(result.logits.shape			,			(self.batch_size, self.num_choices)  )




      def 				lowercase__    (			self			: str  )				->      Optional[Any]:
            _lowerCAmelCase							     =							self.prepare_config_and_inputs()
            ((_lowerCAmelCase)     ,							(_lowerCAmelCase)     ,							(_lowerCAmelCase)     ,							(_lowerCAmelCase)     ,							(_lowerCAmelCase)     ,							(_lowerCAmelCase))							     =							config_and_inputs
            _lowerCAmelCase							     =							{"""input_ids""": input_ids, """attention_mask""": input_mask}
            return config, inputs_dict




@require_torch
class  UpperCAmelCase     (			snake_case_							,		snake_case_							,		unittest.TestCase   ):
      _lowercase:			str						     =   (
          (
              DistilBertModel,
              DistilBertForMaskedLM,
              DistilBertForMultipleChoice,
              DistilBertForQuestionAnswering,
              DistilBertForSequenceClassification,
              DistilBertForTokenClassification,
          )
          if is_torch_available()
          else None
      )
      _lowercase:			Dict						     =   (
          {
              '''feature-extraction''': DistilBertModel,
              '''fill-mask''': DistilBertForMaskedLM,
              '''question-answering''': DistilBertForQuestionAnswering,
              '''text-classification''': DistilBertForSequenceClassification,
              '''token-classification''': DistilBertForTokenClassification,
              '''zero-shot''': DistilBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase:			Any						     =   True
      _lowercase:			int						     =   True
      _lowercase:			Union[str, Any]						     =   True
      _lowercase:			List[Any]						     =   True




      def 				lowercase__    (			self			: Dict  )				->      Optional[Any]:
            _lowerCAmelCase							     =							DistilBertModelTester(self  )
            _lowerCAmelCase							     =							ConfigTester(self			,			config_class=__snake_case			,			dim=37  )




      def 				lowercase__    (			self			: List[str]  )				->      Optional[Any]:
            self.config_tester.run_common_tests()




      def 				lowercase__    (			self			: str  )				->      Any:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_model(*__snake_case  )




      def 				lowercase__    (			self			: Tuple  )				->      Tuple:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_masked_lm(*__snake_case  )




      def 				lowercase__    (			self			: List[Any]  )				->      str:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_question_answering(*__snake_case  )




      def 				lowercase__    (			self			: str  )				->      List[str]:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_sequence_classification(*__snake_case  )




      def 				lowercase__    (			self			: int  )				->      Optional[Any]:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_token_classification(*__snake_case  )




      def 				lowercase__    (			self			: Any  )				->      Union[str, Any]:
            _lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_multiple_choice(*__snake_case  )




      @slow
      def 				lowercase__    (			self			: List[str]  )				->      List[Any]:
            for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  _lowerCAmelCase							     =							DistilBertModel.from_pretrained(__snake_case  )
                  self.assertIsNotNone(__snake_case  )




      @slow
      @require_torch_gpu
      def 				lowercase__    (			self			: List[Any]  )				->      Tuple:
            _lowerCAmelCase     ,							_lowerCAmelCase							     =							self.model_tester.prepare_config_and_inputs_for_common()
            for model_class in self.all_model_classes:
                  # BertForMultipleChoice behaves incorrectly in JIT environments.
                  if model_class == DistilBertForMultipleChoice:
                        return

                  _lowerCAmelCase							     =							True
                  _lowerCAmelCase							     =							model_class(config=__snake_case  )

                  _lowerCAmelCase							     =							self._prepare_for_class(__snake_case			,			__snake_case  )
                  _lowerCAmelCase							     =							torch.jit.trace(
                      __snake_case			,			(inputs_dict["""input_ids"""].to("""cpu"""  ), inputs_dict["""attention_mask"""].to("""cpu"""  ))  )

                  with tempfile.TemporaryDirectory() as tmp:
                        torch.jit.save(__snake_case			,			os.path.join(__snake_case			,			"""traced_model.pt"""  )  )
                        _lowerCAmelCase							     =							torch.jit.load(os.path.join(__snake_case			,			"""traced_model.pt"""  )			,			map_location=__snake_case  )
                        loaded(inputs_dict["""input_ids"""].to(__snake_case  )			,			inputs_dict["""attention_mask"""].to(__snake_case  )  )




@require_torch
class  UpperCAmelCase     (			unittest.TestCase   ):




      @slow
      def 				lowercase__    (			self			: Any  )				->      List[str]:
            _lowerCAmelCase							     =							DistilBertModel.from_pretrained("""distilbert-base-uncased"""  )
            _lowerCAmelCase							     =							torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]  )
            _lowerCAmelCase							     =							torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]  )
            with torch.no_grad():
                  _lowerCAmelCase							     =							model(__snake_case			,			attention_mask=__snake_case  )[0]
            _lowerCAmelCase							     =							torch.Size((1, 11, 7_68)  )
            self.assertEqual(output.shape			,			__snake_case  )
            _lowerCAmelCase							     =							torch.tensor(
                [[[-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(torch.allclose(output[:, 1:4, 1:4]			,			__snake_case			,			atol=1E-4  )  )







'''simple docstring'''





import inspect
import unittest

import numpy as np

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
    import jax

    from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel




class        __A      ( unittest.TestCase	):



    def __init__(self    :    str   ,			__a    :    Optional[Any]   ,			__a    :    Optional[Any]=13   ,			__a    :    int=30   ,			__a    :    Union[str, Any]=2   ,			__a    :    Dict=3   ,			__a    :    List[Any]=True   ,			__a    :    Optional[Any]=True   ,			__a    :    List[Any]=32   ,			__a    :    Any=5   ,			__a    :    str=4   ,			__a    :    Optional[int]=37   ,			__a    :    Optional[int]="gelu"   ,			__a    :    List[str]=0.1   ,			__a    :    Tuple=0.1   ,			__a    :    List[str]=10   ,			__a    :    Optional[int]=0.02   ,			):
         UpperCAmelCase_           =   parent
         UpperCAmelCase_           =   batch_size
         UpperCAmelCase_           =   image_size
         UpperCAmelCase_           =   patch_size
         UpperCAmelCase_           =   num_channels
         UpperCAmelCase_           =   is_training
         UpperCAmelCase_           =   use_labels
         UpperCAmelCase_           =   hidden_size
         UpperCAmelCase_           =   num_hidden_layers
         UpperCAmelCase_           =   num_attention_heads
         UpperCAmelCase_           =   intermediate_size
         UpperCAmelCase_           =   hidden_act
         UpperCAmelCase_           =   hidden_dropout_prob
         UpperCAmelCase_           =   attention_probs_dropout_prob
         UpperCAmelCase_           =   type_sequence_label_size
         UpperCAmelCase_           =   initializer_range

         # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
         UpperCAmelCase_           =   (image_size // patch_size) ** 2
         UpperCAmelCase_           =   num_patches + 1



    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

         UpperCAmelCase_           =   ViTConfig(
             image_size=self.image_size   ,			patch_size=self.patch_size   ,			num_channels=self.num_channels   ,			hidden_size=self.hidden_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   ,			is_decoder=__a   ,			initializer_range=self.initializer_range   ,			)

         return config, pixel_values



    def 					_lowercase						(self    :    Dict   ,			__a    :    Any   ,			__a    :    List[Any] ):
         UpperCAmelCase_           =   FlaxViTModel(config=__a )
         UpperCAmelCase_           =   model(__a )
         # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
         UpperCAmelCase_           =   (self.image_size, self.image_size)
         UpperCAmelCase_           =   (self.patch_size, self.patch_size)
         UpperCAmelCase_           =   (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
         self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, num_patches + 1, self.hidden_size) )



    def 					_lowercase						(self    :    Tuple   ,			__a    :    str   ,			__a    :    Any ):
         UpperCAmelCase_           =   self.type_sequence_label_size
         UpperCAmelCase_           =   FlaxViTForImageClassification(config=__a )
         UpperCAmelCase_           =   model(__a )
         self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.type_sequence_label_size) )

         # test greyscale images
         UpperCAmelCase_           =   1
         UpperCAmelCase_           =   FlaxViTForImageClassification(__a )

         UpperCAmelCase_           =   floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
         UpperCAmelCase_           =   model(__a )






    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_           =   self.prepare_config_and_inputs()
         (
             (
             UpperCAmelCase_
         )    , (
             UpperCAmelCase_
         )    , 
         )           =   config_and_inputs
         UpperCAmelCase_           =   {"pixel_values": pixel_values}
         return config, inputs_dict




@require_flax
class        __A      ( UpperCamelCase__		,	unittest.TestCase	):
    a__       :						Tuple          =  (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()



    def 					_lowercase						(self    :    Any ):
         UpperCAmelCase_           =   FlaxViTModelTester(self )
         UpperCAmelCase_           =   ConfigTester(self   ,			config_class=__a   ,			has_text_modality=__a   ,			hidden_size=37 )



    def 					_lowercase						(self    :    Tuple ):
         self.config_tester.run_common_tests()



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__a )



    def 					_lowercase						(self    :    str ):
         UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_image_classification(*__a )



    def 					_lowercase						(self    :    Tuple ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              UpperCAmelCase_           =   model_class(__a )
              UpperCAmelCase_           =   inspect.signature(model.__call__ )
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              UpperCAmelCase_           =   [*signature.parameters.keys()]

              UpperCAmelCase_           =   ["pixel_values"]
              self.assertListEqual(arg_names[:1]   ,			__a )



    def 					_lowercase						(self    :    Optional[Any] ):
         UpperCAmelCase_    , UpperCAmelCase_           =   self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__ ):
                   UpperCAmelCase_           =   self._prepare_for_class(__a   ,			__a )
                   UpperCAmelCase_           =   model_class(__a )

                   @jax.jit
                   def model_jitted(__a    :    Tuple   ,			**__a    :    List[Any] ):
                        return model(pixel_values=__a   ,			**__a )

                   with self.subTest("JIT Enabled" ):
                        UpperCAmelCase_           =   model_jitted(**__a ).to_tuple()

                   with self.subTest("JIT Disabled" ):
                        with jax.disable_jit():
                             UpperCAmelCase_           =   model_jitted(**__a ).to_tuple()

                   self.assertEqual(len(__a )   ,			len(__a ) )
                   for jitted_output, output in zip(__a   ,			__a ):
                        self.assertEqual(jitted_output.shape   ,			output.shape )






    @slow
    def 					_lowercase						(self    :    Tuple ):
         for model_class_name in self.all_model_classes:
              UpperCAmelCase_           =   model_class_name.from_pretrained("google/vit-base-patch16-224" )
              UpperCAmelCase_           =   model(np.ones((1, 3, 224, 224) ) )
              self.assertIsNotNone(__a )




