


import unittest

from transformers import AutoConfig, AutoTokenizer, BertConfig, TensorType, is_flax_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, require_flax, slow


if is_flax_available():
   import jax

   from transformers.models.auto.modeling_flax_auto import FlaxAutoModel
   from transformers.models.bert.modeling_flax_bert import FlaxBertModel
   from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel
@require_flax
class 				SCREAMING_SNAKE_CASE__						(       unittest.TestCase    ):
      @slow
      def  SCREAMING_SNAKE_CASE							( self       : int  )	->    Tuple:
             for model_name in ["bert-base-cased", "bert-large-uncased"]:
                    with self.subTest(SCREAMING_SNAKE_CASE__  ):
                           a_      :	Any							=						AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__  )
                           self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )
                           self.assertIsInstance(SCREAMING_SNAKE_CASE__  ,							SCREAMING_SNAKE_CASE__  )

                           a_      :	List[Any]							=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
                           self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )
                           self.assertIsInstance(SCREAMING_SNAKE_CASE__  ,							SCREAMING_SNAKE_CASE__  )
      @slow
      def  SCREAMING_SNAKE_CASE							( self       : str  )	->    Union[str, Any]:
             for model_name in ["roberta-base", "roberta-large"]:
                    with self.subTest(SCREAMING_SNAKE_CASE__  ):
                           a_      :	List[Any]							=						AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE__  )
                           self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )
                           self.assertIsInstance(SCREAMING_SNAKE_CASE__  ,							SCREAMING_SNAKE_CASE__  )

                           a_      :	Optional[int]							=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
                           self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )
                           self.assertIsInstance(SCREAMING_SNAKE_CASE__  ,							SCREAMING_SNAKE_CASE__  )
      @slow
      def  SCREAMING_SNAKE_CASE							( self       : int  )	->    int:
             for model_name in ["bert-base-cased", "bert-large-uncased"]:
                    a_      :	List[str]							=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__  )
                    a_      :	Optional[Any]							=						FlaxBertModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
                    a_      :	Optional[int]							=						tokenizer('Do you support jax jitted function?'  ,							return_tensors=TensorType.JAX  )

                    @jax.jit
                    def eval(**SCREAMING_SNAKE_CASE__       : Any  ):
                           return model(**SCREAMING_SNAKE_CASE__  )

                    eval(**SCREAMING_SNAKE_CASE__  ).block_until_ready()
      @slow
      def  SCREAMING_SNAKE_CASE							( self       : Dict  )	->    Dict:
             for model_name in ["roberta-base", "roberta-large"]:
                    a_      :	Union[str, Any]							=						AutoTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__  )
                    a_      :	int							=						FlaxRobertaModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
                    a_      :	Tuple							=						tokenizer('Do you support jax jitted function?'  ,							return_tensors=TensorType.JAX  )

                    @jax.jit
                    def eval(**SCREAMING_SNAKE_CASE__       : List[str]  ):
                           return model(**SCREAMING_SNAKE_CASE__  )

                    eval(**SCREAMING_SNAKE_CASE__  ).block_until_ready()
      def  SCREAMING_SNAKE_CASE							( self       : Optional[int]  )	->    List[Any]:
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE__  ,							'bert-base is not a local folder and is not a valid model identifier'  ):
                    a_      :	Optional[int]							=						FlaxAutoModel.from_pretrained('bert-base'  )
      def  SCREAMING_SNAKE_CASE							( self       : int  )	->    Optional[Any]:
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE__  ,							r'aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'  ):
                    a_      :	List[str]							=						FlaxAutoModel.from_pretrained(SCREAMING_SNAKE_CASE__  ,							revision='aaaaaa'  )
      def  SCREAMING_SNAKE_CASE							( self       : Any  )	->    Dict:
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE__  ,							'hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack'  ,							):
                    a_      :	Any							=						FlaxAutoModel.from_pretrained('hf-internal-testing/config-no-model'  )

      def  SCREAMING_SNAKE_CASE							( self       : List[str]  )	->    Tuple:
             with self.assertRaisesRegex(SCREAMING_SNAKE_CASE__  ,							'Use `from_pt=True` to load this model'  ):
                    a_      :	Optional[Any]							=						FlaxAutoModel.from_pretrained('hf-internal-testing/tiny-bert-pt-only'  )






UpperCAmelCase_   :				Optional[int]						      = '\n# Transformers installation\n! pip install transformers datasets\n# To install from source instead of the last release, comment the command above and uncomment the following one.\n# ! pip install git+https://github.com/huggingface/transformers.git\n'

UpperCAmelCase_   :				List[str]						      = [{'type': 'code', 'content': INSTALL_CONTENT}]
UpperCAmelCase_   :				int						      = {
    '{processor_class}': 'FakeProcessorClass',
    '{model_class}': 'FakeModelClass',
    '{object_class}': 'FakeObjectClass',
}


