





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



					'''simple docstring'''

					@slow
					def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[int]	)					->		Optional[int]:




											'''simple docstring'''


											for model_name in ["bert-base-cased", "bert-large-uncased"]:
																	with self.subTest(lowerCamelCase__	):
																							SCREAMING_SNAKE_CASE		=						AutoConfig.from_pretrained(lowerCamelCase__	)
																							self.assertIsNotNone(lowerCamelCase__	)
																							self.assertIsInstance(lowerCamelCase__						,lowerCamelCase__	)

																							SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained(lowerCamelCase__	)
																							self.assertIsNotNone(lowerCamelCase__	)
																							self.assertIsInstance(lowerCamelCase__						,lowerCamelCase__	)






					@slow
					def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		List[str]:




											'''simple docstring'''


											for model_name in ["roberta-base", "roberta-large"]:
																	with self.subTest(lowerCamelCase__	):
																							SCREAMING_SNAKE_CASE		=						AutoConfig.from_pretrained(lowerCamelCase__	)
																							self.assertIsNotNone(lowerCamelCase__	)
																							self.assertIsInstance(lowerCamelCase__						,lowerCamelCase__	)

																							SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained(lowerCamelCase__	)
																							self.assertIsNotNone(lowerCamelCase__	)
																							self.assertIsInstance(lowerCamelCase__						,lowerCamelCase__	)






					@slow
					def     SCREAMING_SNAKE_CASE__  (							self						:   Tuple	)					->		int:




											'''simple docstring'''


											for model_name in ["bert-base-cased", "bert-large-uncased"]:
																	SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained(lowerCamelCase__	)
																	SCREAMING_SNAKE_CASE		=						FlaxBertModel.from_pretrained(lowerCamelCase__	)
																	SCREAMING_SNAKE_CASE		=						tokenizer("""Do you support jax jitted function?"""						,return_tensors=TensorType.JAX	)

																	@jax.jit
																	def eval(**lowerCamelCase__						:   Union[str, Any]	):
																							return model(**lowerCamelCase__	)

																	eval(**lowerCamelCase__	).block_until_ready()






					@slow
					def     SCREAMING_SNAKE_CASE__  (							self						:   List[Any]	)					->		int:




											'''simple docstring'''


											for model_name in ["roberta-base", "roberta-large"]:
																	SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained(lowerCamelCase__	)
																	SCREAMING_SNAKE_CASE		=						FlaxRobertaModel.from_pretrained(lowerCamelCase__	)
																	SCREAMING_SNAKE_CASE		=						tokenizer("""Do you support jax jitted function?"""						,return_tensors=TensorType.JAX	)

																	@jax.jit
																	def eval(**lowerCamelCase__						:   Union[str, Any]	):
																							return model(**lowerCamelCase__	)

																	eval(**lowerCamelCase__	).block_until_ready()






					def     SCREAMING_SNAKE_CASE__  (							self						:   List[str]	)					->		Dict:




											'''simple docstring'''


											with self.assertRaisesRegex(
											    lowerCamelCase__						,"""bert-base is not a local folder and is not a valid model identifier"""	):
																	SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained("""bert-base"""	)






					def     SCREAMING_SNAKE_CASE__  (							self						:   str	)					->		Optional[Any]:




											'''simple docstring'''


											with self.assertRaisesRegex(
											    lowerCamelCase__						,R"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""	):
																	SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained(lowerCamelCase__						,revision="""aaaaaa"""	)






					def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		int:




											'''simple docstring'''


											with self.assertRaisesRegex(
											    lowerCamelCase__						,"""hf-internal-testing/config-no-model does not appear to have a file named flax_model.msgpack"""						,):
																	SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""	)






					def     SCREAMING_SNAKE_CASE__  (							self						:   Any	)					->		List[str]:




											'''simple docstring'''


											with self.assertRaisesRegex(lowerCamelCase__						,"""Use `from_pt=True` to load this model"""	):
																	SCREAMING_SNAKE_CASE		=						FlaxAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""	)









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


SCREAMING_SNAKE_CASE_						       =	{
    """configuration_megatron_bert""": ["""MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MegatronBertConfig"""],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      SCREAMING_SNAKE_CASE_						       =	[
          """MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MegatronBertForCausalLM""",
          """MegatronBertForMaskedLM""",
          """MegatronBertForMultipleChoice""",
          """MegatronBertForNextSentencePrediction""",
          """MegatronBertForPreTraining""",
          """MegatronBertForQuestionAnswering""",
          """MegatronBertForSequenceClassification""",
          """MegatronBertForTokenClassification""",
          """MegatronBertModel""",
          """MegatronBertPreTrainedModel""",
      ]

if TYPE_CHECKING:
      from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_megatron_bert import (
                MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                MegatronBertForCausalLM,
                MegatronBertForMaskedLM,
                MegatronBertForMultipleChoice,
                MegatronBertForNextSentencePrediction,
                MegatronBertForPreTraining,
                MegatronBertForQuestionAnswering,
                MegatronBertForSequenceClassification,
                MegatronBertForTokenClassification,
                MegatronBertModel,
                MegatronBertPreTrainedModel,
            )

else:
      import sys

      SCREAMING_SNAKE_CASE_						       =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


