






import inspect
import unittest

from transformers import RegNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

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


if is_flax_available():
     import jax
     import jax.numpy as jnp

     from transformers.models.regnet.modeling_flax_regnet import FlaxRegNetForImageClassification, FlaxRegNetModel

if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor




class       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       def __init__(				self      ,    a__      ,    a__=3      ,    a__=32      ,    a__=3      ,    a__=10      ,    a__=[10, 20, 30, 40]      ,    a__=[1, 1, 2, 1]      ,    a__=True      ,    a__=True      ,    a__="relu"      ,    a__=3      ,    a__=None      ,    )			->   Any:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    parent
            __snake_case      :Optional[Any] 							=    batch_size
            __snake_case      :int 							=    image_size
            __snake_case      :Dict 							=    num_channels
            __snake_case      :Any 							=    embeddings_size
            __snake_case      :Dict 							=    hidden_sizes
            __snake_case      :Union[str, Any] 							=    depths
            __snake_case      :List[str] 							=    is_training
            __snake_case      :List[Any] 							=    use_labels
            __snake_case      :int 							=    hidden_act
            __snake_case      :List[str] 							=    num_labels
            __snake_case      :Union[str, Any] 							=    scope
            __snake_case      :str 							=    len(lowerCamelCase_     )



       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            __snake_case      :Optional[int] 							=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

            __snake_case      :Dict 							=    self.get_config()

            return config, pixel_values



       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            return RegNetConfig(
                num_channels=self.num_channels      ,    embeddings_size=self.embeddings_size      ,    hidden_sizes=self.hidden_sizes      ,    depths=self.depths      ,    hidden_act=self.hidden_act      ,    num_labels=self.num_labels      ,    image_size=self.image_size      ,    )



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Optional[Any] 							=    FlaxRegNetModel(config=lowerCamelCase_     )
            __snake_case      :Dict 							=    model(lowerCamelCase_     )

            # Output shape (b, c, h, w)
            self.parent.assertEqual(
                result.last_hidden_state.shape      ,    (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)      ,    )



       def        __lowercase    (				self      ,    a__      ,    a__     )			->   Any:






            '''simple docstring'''





            __snake_case      :int 							=    self.num_labels
            __snake_case      :Optional[Any] 							=    FlaxRegNetForImageClassification(config=lowerCamelCase_     )
            __snake_case      :Any 							=    model(lowerCamelCase_     )
            self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.num_labels)     )



       def        __lowercase    (				self     )			->   int:






            '''simple docstring'''





            __snake_case      :Tuple 							=    self.prepare_config_and_inputs()
            __snake_case							,			__snake_case      :Optional[int] 							=    config_and_inputs
            __snake_case      :Tuple 							=    {"""pixel_values""": pixel_values}
            return config, inputs_dict




@require_flax
class       snake_case__ (				a__						,     unittest.TestCase):

       '''simple docstring'''
       lowerCamelCase				:					Any								=			(FlaxRegNetModel, FlaxRegNetForImageClassification) if is_flax_available() else ()

       lowerCamelCase				:					List[Any]								=			False
       lowerCamelCase				:					Union[str, Any]								=			False
       lowerCamelCase				:					Optional[int]								=			False



       def        __lowercase    (				self     )			->   None:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    FlaxRegNetModelTester(self     )
            __snake_case      :List[Any] 							=    ConfigTester(self      ,    config_class=lowerCamelCase_      ,    has_text_modality=lowerCamelCase_     )



       def        __lowercase    (				self     )			->   Dict:






            '''simple docstring'''





            self.create_and_test_config_common_properties()
            self.config_tester.create_and_test_config_to_json_string()
            self.config_tester.create_and_test_config_to_json_file()
            self.config_tester.create_and_test_config_from_and_save_pretrained()
            self.config_tester.create_and_test_config_with_num_labels()
            self.config_tester.check_config_can_be_init_without_params()
            self.config_tester.check_config_arguments_init()



       def        __lowercase    (				self     )			->   List[Any]:






            '''simple docstring'''





            return



       def        __lowercase    (				self     )			->   Optional[Any]:






            '''simple docstring'''





            __snake_case      :Tuple 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*lowerCamelCase_     )



       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            __snake_case      :Dict 							=    self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*lowerCamelCase_     )



       @unittest.skip(reason="""RegNet does not use inputs_embeds"""     )
       def        __lowercase    (				self     )			->   List[str]:






            '''simple docstring'''





            pass



       @unittest.skip(reason="""RegNet does not support input and output embeddings"""     )
       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            pass



       def        __lowercase    (				self     )			->   Tuple:






            '''simple docstring'''





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

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

                 __snake_case      :Dict 							=    ["""pixel_values"""]
                 self.assertListEqual(arg_names[:1]      ,    lowerCamelCase_     )



       def        __lowercase    (				self     )			->   Any:






            '''simple docstring'''





            def check_hidden_states_output(a__      ,    a__      ,    a__     ):
                 __snake_case      :Optional[int] 							=    model_class(lowerCamelCase_     )

                 __snake_case      :int 							=    model(**self._prepare_for_class(lowerCamelCase_      ,    lowerCamelCase_     )     )

                 __snake_case      :str 							=    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                 __snake_case      :Any 							=    self.model_tester.num_stages
                 self.assertEqual(len(lowerCamelCase_     )      ,    expected_num_stages + 1     )

            __snake_case							,			__snake_case      :Union[str, Any] 							=    self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 __snake_case      :Optional[Any] 							=    True
                 check_hidden_states_output(lowerCamelCase_      ,    lowerCamelCase_      ,    lowerCamelCase_     )

                 # check that output_hidden_states also work using config
                 del inputs_dict["output_hidden_states"]
                 __snake_case      :Dict 							=    True

                 check_hidden_states_output(lowerCamelCase_      ,    lowerCamelCase_      ,    lowerCamelCase_     )



       def        __lowercase    (				self     )			->   Union[str, Any]:






            '''simple docstring'''





            __snake_case							,			__snake_case      :Tuple 							=    self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                 with self.subTest(model_class.__name__     ):
                      __snake_case      :Any 							=    self._prepare_for_class(lowerCamelCase_      ,    lowerCamelCase_     )
                      __snake_case      :Dict 							=    model_class(lowerCamelCase_     )

                      @jax.jit
                      def model_jitted(a__      ,    **a__     ):
                           return model(pixel_values=lowerCamelCase_      ,    **lowerCamelCase_     )

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

                      with self.subTest("""JIT Disabled"""     ):
                           with jax.disable_jit():
                                __snake_case      :Optional[Any] 							=    model_jitted(**lowerCamelCase_     ).to_tuple()

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






def 	UpperCamelCase	(							):




     '''simple docstring'''






     __snake_case      :Union[str, Any] 							=    Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""	)
     return image




@require_flax
class       snake_case__ (				unittest.TestCase):

       '''simple docstring'''



       @cached_property
       def        __lowercase    (				self     )			->   Optional[int]:






            '''simple docstring'''





            return AutoImageProcessor.from_pretrained("""facebook/regnet-y-040"""     ) if is_vision_available() else None



       @slow
       def        __lowercase    (				self     )			->   str:






            '''simple docstring'''





            __snake_case      :List[Any] 							=    FlaxRegNetForImageClassification.from_pretrained("""facebook/regnet-y-040"""     )

            __snake_case      :Optional[int] 							=    self.default_image_processor
            __snake_case      :Optional[int] 							=    prepare_img()
            __snake_case      :Tuple 							=    image_processor(images=lowerCamelCase_      ,    return_tensors="""np"""     )

            __snake_case      :List[str] 							=    model(**lowerCamelCase_     )

            # verify the logits
            __snake_case      :List[str] 							=    (1, 10_00)
            self.assertEqual(outputs.logits.shape      ,    lowerCamelCase_     )

            __snake_case      :Optional[Any] 							=    jnp.array([-0.41_80, -1.50_51, -3.48_36]     )

            self.assertTrue(jnp.allclose(outputs.logits[0, :3]      ,    lowerCamelCase_      ,    atol=1e-4     )     )






'''simple docstring'''






import os
from typing import List, Optional, Union

from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_utils_base import AddedToken
from ...utils import logging


__UpperCAmelCase       							=				logging.get_logger(__name__)

__UpperCAmelCase       							=				{'''vocab_file''': '''vocab.txt'''}

__UpperCAmelCase       							=				{
    '''vocab_file''': {
        '''facebook/esm2_t6_8M_UR50D''': '''https://huggingface.co/facebook/esm2_t6_8M_UR50D/resolve/main/vocab.txt''',
        '''facebook/esm2_t12_35M_UR50D''': '''https://huggingface.co/facebook/esm2_t12_35M_UR50D/resolve/main/vocab.txt''',
    },
}

__UpperCAmelCase       							=				{
    '''facebook/esm2_t6_8M_UR50D''': 1_024,
    '''facebook/esm2_t12_35M_UR50D''': 1_024,
}


def    _snake_case  (	A			)		->       Optional[Any]:
							with open(A					,       '''r'''			) as f:
														lowerCAmelCase__			  =      f.read().splitlines()
														return [l.strip() for l in lines]






class 						a__							(				a__    ):





							'''simple docstring'''






							lowercase__     :  Optional[Any]      		= VOCAB_FILES_NAMES
							lowercase__     :  List[Any]      		= PRETRAINED_VOCAB_FILES_MAP
							lowercase__     :  Union[str, Any]      		= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowercase__     :  Union[str, Any]      		= ["input_ids", "attention_mask"]







							def __init__(						self   ,       lowerCamelCase_   ,       lowerCamelCase_="<unk>"   ,       lowerCamelCase_="<cls>"   ,       lowerCamelCase_="<pad>"   ,       lowerCamelCase_="<mask>"   ,       lowerCamelCase_="<eos>"   ,       **lowerCamelCase_   ,       )			->				Tuple:
														super().__init__(**lowerCamelCase_  )
														lowerCAmelCase__			  =      load_vocab_file(lowerCamelCase_  )
														lowerCAmelCase__			  =      dict(enumerate(self.all_tokens  )  )
														lowerCAmelCase__			  =      {tok: ind for ind, tok in enumerate(self.all_tokens  )}
														lowerCAmelCase__			  =      unk_token
														lowerCAmelCase__			  =      cls_token
														lowerCAmelCase__			  =      pad_token
														lowerCAmelCase__			  =      mask_token
														lowerCAmelCase__			  =      eos_token
														lowerCAmelCase__			  =      self.all_tokens
														self._create_trie(self.unique_no_split_tokens  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				str:
														return self._id_to_token.get(lowerCamelCase_   ,       self.unk_token  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				int:
														return self._token_to_id.get(lowerCamelCase_   ,       self._token_to_id.get(self.unk_token  )  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       **lowerCamelCase_  )			->				Union[str, Any]:
														return text.split()







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_=False  )			->				Dict:
														return len(self._id_to_token  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return {token: i for i, token in enumerate(self.all_tokens  )}







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				int:
														return self._token_to_id.get(lowerCamelCase_   ,       self._token_to_id.get(self.unk_token  )  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				str:
														return self._id_to_token.get(lowerCamelCase_   ,       self.unk_token  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = None  )			->				List[int]:
														lowerCAmelCase__			  =      [self.cls_token_id]
														lowerCAmelCase__			  =      [self.eos_token_id]  # No sep token in ESM vocabulary
														if token_ids_a is None:
																					if self.eos_token_id is None:
																												return cls + token_ids_a
																					else:
																												return cls + token_ids_a + sep
														elif self.eos_token_id is None:
																					raise ValueError('''Cannot tokenize multiple sequences when EOS token is not set!'''  )
														return cls + token_ids_a + sep + token_ids_a + sep  # Multiple inputs always have an EOS token







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = None   ,       lowerCamelCase_ = False  )			->				List[int]:
														if already_has_special_tokens:
																					if token_ids_a is not None:
																												raise ValueError(
																												    '''You should not supply a second sequence if the provided sequence of '''
																												    '''ids is already formatted with special tokens for the model.'''  )

																					return [1 if token in self.all_special_ids else 0 for token in token_ids_a]
														lowerCAmelCase__			  =      [1] + ([0] * len(lowerCamelCase_  )) + [1]
														if token_ids_a is not None:
																					mask += [0] * len(lowerCamelCase_  ) + [1]
														return mask







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Union[str, Any]:
														lowerCAmelCase__			  =      os.path.join(lowerCamelCase_   ,       (filename_prefix + '''-''' if filename_prefix else '''''') + '''vocab.txt'''  )
														with open(lowerCamelCase_   ,       '''w'''  ) as f:
																					f.write('''\n'''.join(self.all_tokens  )  )
														return (vocab_file,)







							@property
							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														return self.get_vocab_size(with_added_tokens=lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = False  )			->				int:
														return super()._add_tokens(lowerCamelCase_   ,       special_tokens=lowerCamelCase_  )