


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


  '''simple docstring'''







  def __init__(							self :       Dict	,_a :       Union[str, Any]	,_a :       Optional[Any]=13	,_a :       List[str]=30	,_a :       List[str]=2	,_a :       Dict=3	,_a :       Tuple=True	,_a :       Optional[int]=True	,_a :       Any=32	,_a :       List[str]=5	,_a :       Optional[int]=4	,_a :       List[Any]=37	,_a :       int="gelu"	,_a :       int=0.1	,_a :       Optional[int]=0.1	,_a :       Union[str, Any]=10	,_a :       Tuple=0.02	,):

        '''simple docstring'''




        A_						:  Union[str, Any]          =				parent
        A_						:  Optional[Any]          =				batch_size
        A_						:  Optional[int]          =				image_size
        A_						:  List[Any]          =				patch_size
        A_						:  Optional[int]          =				num_channels
        A_						:  List[str]          =				is_training
        A_						:  Union[str, Any]          =				use_labels
        A_						:  int          =				hidden_size
        A_						:  Optional[int]          =				num_hidden_layers
        A_						:  List[Any]          =				num_attention_heads
        A_						:  Union[str, Any]          =				intermediate_size
        A_						:  Optional[int]          =				hidden_act
        A_						:  Tuple          =				hidden_dropout_prob
        A_						:  Optional[Any]          =				attention_probs_dropout_prob
        A_						:  int          =				type_sequence_label_size
        A_						:  List[str]          =				initializer_range

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






  def        _a							(							self :       Tuple ):

        '''simple docstring'''




        A_						:  Optional[Any]          =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

        A_						:  str          =				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        _a							(							self :       Optional[int]	,_a :       Tuple	,_a :       Tuple ):

        '''simple docstring'''




        A_						:  Any          =				FlaxViTModel(config=_a )
        A_						:  str          =				model(_a )
        # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
        A_						:  str          =				(self.image_size, self.image_size)
        A_						:  List[str]          =				(self.patch_size, self.patch_size)
        A_						:  Optional[int]          =				(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        _a							(							self :       List[Any]	,_a :       Any	,_a :       Union[str, Any] ):

        '''simple docstring'''




        A_						:  List[Any]          =				self.type_sequence_label_size
        A_						:  Tuple          =				FlaxViTForImageClassification(config=_a )
        A_						:  int          =				model(_a )
        self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.type_sequence_label_size) )

        # test greyscale images
        A_						:  List[str]          =				1
        A_						:  str          =				FlaxViTForImageClassification(_a )

        A_						:  int          =				floats_tensor([self.batch_size, 1, self.image_size, self.image_size] )
        A_						:  Any          =				model(_a )






  def        _a							(							self :       int ):

        '''simple docstring'''




        A_						:  List[str]          =				self.prepare_config_and_inputs()
        (
            (
            A_
        )   ,					(
            A_
        )   ,					
        )						:  int          =				config_and_inputs
        A_						:  str          =				{"""pixel_values""": pixel_values}
        return config, inputs_dict







@require_flax
class 	__lowerCAmelCase		(							__SCREAMING_SNAKE_CASE  , unittest.TestCase    ):


  '''simple docstring'''







  a_  					=							(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()






  def        _a							(							self :       Dict ):

        '''simple docstring'''




        A_						:  List[str]          =				FlaxViTModelTester(self )
        A_						:  str          =				ConfigTester(self	,config_class=_a	,has_text_modality=_a	,hidden_size=37 )






  def        _a							(							self :       Optional[Any] ):

        '''simple docstring'''




        self.config_tester.run_common_tests()






  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        A_						:  Any          =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_a )






  def        _a							(							self :       Optional[Any] ):

        '''simple docstring'''




        A_						:  int          =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*_a )






  def        _a							(							self :       Optional[Any] ):

        '''simple docstring'''




        A_   ,					A_						:  str          =				self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              A_						:  List[Any]          =				model_class(_a )
              A_						:  Optional[int]          =				inspect.signature(model.__call__ )
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              A_						:  List[str]          =				[*signature.parameters.keys()]

              A_						:  Tuple          =				["""pixel_values"""]
              self.assertListEqual(arg_names[:1]	,_a )






  def        _a							(							self :       List[str] ):

        '''simple docstring'''




        A_   ,					A_						:  Optional[int]          =				self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__ ):
                    A_						:  List[str]          =				self._prepare_for_class(_a	,_a )
                    A_						:  Tuple          =				model_class(_a )

                    @jax.jit
                    def model_jitted(_a :       Dict	,**_a :       Dict ):
                          return model(pixel_values=_a	,**_a )

                    with self.subTest("""JIT Enabled""" ):
                          A_						:  List[str]          =				model_jitted(**_a ).to_tuple()

                    with self.subTest("""JIT Disabled""" ):
                          with jax.disable_jit():
                                A_						:  Union[str, Any]          =				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        _a							(							self :       List[str] ):

        '''simple docstring'''




        for model_class_name in self.all_model_classes:
              A_						:  int          =				model_class_name.from_pretrained("""google/vit-base-patch16-224""" )
              A_						:  str          =				model(np.ones((1, 3, 224, 224) ) )
              self.assertIsNotNone(_a )




'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__magic_name__			=      {}

try:
     if not is_sentencepiece_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__			=      ['NllbTokenizer']

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     __magic_name__			=      ['NllbTokenizerFast']


if TYPE_CHECKING:
     try:
          if not is_sentencepiece_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_nllb import NllbTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_nllb_fast import NllbTokenizerFast

else:
     import sys

     __magic_name__			=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
