


'''simple docstring'''







from __future__ import annotations

import copy
import tempfile
import unittest

from transformers import CONFIG_MAPPING, AutoConfig, BertConfig, GPTaConfig, TaConfig, TapasConfig, is_tf_available
from transformers.testing_utils import (
    DUMMY_UNKNOWN_IDENTIFIER,
    SMALL_MODEL_IDENTIFIER,
    RequestCounter,
    require_tensorflow_probability,
    require_tf,
    slow,
)

from ..bert.test_modeling_bert import BertModelTester


if is_tf_available():
   from transformers import (
       TFAutoModel,
       TFAutoModelForCausalLM,
       TFAutoModelForMaskedLM,
       TFAutoModelForPreTraining,
       TFAutoModelForQuestionAnswering,
       TFAutoModelForSeqaSeqLM,
       TFAutoModelForSequenceClassification,
       TFAutoModelForTableQuestionAnswering,
       TFAutoModelForTokenClassification,
       TFAutoModelWithLMHead,
       TFBertForMaskedLM,
       TFBertForPreTraining,
       TFBertForQuestionAnswering,
       TFBertForSequenceClassification,
       TFBertModel,
       TFFunnelBaseModel,
       TFFunnelModel,
       TFGPTaLMHeadModel,
       TFRobertaForMaskedLM,
       TFTaForConditionalGeneration,
       TFTapasForQuestionAnswering,
   )
   from transformers.models.auto.modeling_tf_auto import (
       TF_MODEL_FOR_CAUSAL_LM_MAPPING,
       TF_MODEL_FOR_MASKED_LM_MAPPING,
       TF_MODEL_FOR_PRETRAINING_MAPPING,
       TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
       TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
       TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
       TF_MODEL_MAPPING,
   )
   from transformers.models.bert.modeling_tf_bert import TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST
   from transformers.models.gpta.modeling_tf_gpta import TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST
   from transformers.models.ta.modeling_tf_ta import TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST
   from transformers.models.tapas.modeling_tf_tapas import TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST







class A__       (						UpperCAmelCase__      ):
      __UpperCamelCase       :  Any	       =							"new-model"


if is_tf_available():







   class A__       (						UpperCAmelCase__      ):
         __UpperCamelCase       :  Any	       =							NewModelConfig







@require_tf
class A__       (						unittest.TestCase      ):






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




             '''simple docstring'''
             _a    : Union[str, Any]     					="""bert-base-cased"""
             _a    : Optional[int]     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[int]     					=TFAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             _a    : Tuple     					="""bert-base-cased"""
             _a    : Optional[Any]     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             _a    : Dict     					=TFAutoModelForPreTraining.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Any     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : Optional[int]     					=TFAutoModelForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE	)
                    _a					,    _a    : Dict     					=TFAutoModelForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE     ,						output_loading_info=SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Tuple     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : str     					=TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:int	)       ->					int:




             '''simple docstring'''
             for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Optional[Any]     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : Union[str, Any]     					=TFAutoModelForMaskedLM.from_pretrained(SCREAMING_SNAKE_CASE	)
                    _a					,    _a    : Dict     					=TFAutoModelForMaskedLM.from_pretrained(SCREAMING_SNAKE_CASE     ,						output_loading_info=SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Dict     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : List[Any]     					=TFAutoModelForSeqaSeqLM.from_pretrained(SCREAMING_SNAKE_CASE	)
                    _a					,    _a    : int     					=TFAutoModelForSeqaSeqLM.from_pretrained(SCREAMING_SNAKE_CASE     ,						output_loading_info=SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:int	)       ->					Union[str, Any]:




             '''simple docstring'''
             # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             for model_name in ["bert-base-uncased"]:
                    _a    : str     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : int     					=TFAutoModelForSequenceClassification.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:str	)       ->					Tuple:




             '''simple docstring'''
             # for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             for model_name in ["bert-base-uncased"]:
                    _a    : List[str]     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : List[Any]     					=TFAutoModelForQuestionAnswering.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






      @slow
      @require_tensorflow_probability
      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					List[str]:




             '''simple docstring'''
             for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
                    _a    : str     					=AutoConfig.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                    _a    : str     					=TFAutoModelForTableQuestionAnswering.from_pretrained(SCREAMING_SNAKE_CASE	)
                    _a					,    _a    : Dict     					=TFAutoModelForTableQuestionAnswering.from_pretrained(
                        SCREAMING_SNAKE_CASE     ,						output_loading_info=SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Dict	)       ->					List[Any]:




             '''simple docstring'''
             _a    : Union[str, Any]     					=TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
             self.assertEqual(model.num_parameters()     ,						1_4_4_1_0	)
             self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE	)     ,						1_4_4_1_0	)






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					str:




             '''simple docstring'''
             _a    : str     					=TFAutoModelWithLMHead.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
             self.assertEqual(model.num_parameters()     ,						1_4_4_1_0	)
             self.assertEqual(model.num_parameters(only_trainable=SCREAMING_SNAKE_CASE	)     ,						1_4_4_1_0	)






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




             '''simple docstring'''
             # For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
             _a    : Any     					=TFAutoModel.from_pretrained("""sgugger/funnel-random-tiny"""	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             _a    : Optional[Any]     					=copy.deepcopy(model.config	)
             _a    : List[Any]     					=["""FunnelBaseModel"""]
             _a    : List[str]     					=TFAutoModel.from_config(SCREAMING_SNAKE_CASE	)
             self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             with tempfile.TemporaryDirectory() as tmp_dir:
                    model.save_pretrained(SCREAMING_SNAKE_CASE	)
                    _a    : Union[str, Any]     					=TFAutoModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             try:
                    AutoConfig.register("""new-model"""     ,						SCREAMING_SNAKE_CASE	)

                    _a    : Union[str, Any]     					=[
                        TFAutoModel,
                        TFAutoModelForCausalLM,
                        TFAutoModelForMaskedLM,
                        TFAutoModelForPreTraining,
                        TFAutoModelForQuestionAnswering,
                        TFAutoModelForSequenceClassification,
                        TFAutoModelForTokenClassification,
                    ]

                    for auto_class in auto_classes:
                           with self.subTest(auto_class.__name__	):
                                  # Wrong config class will raise an error
                                  with self.assertRaises(SCREAMING_SNAKE_CASE	):
                                         auto_class.register(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
                                  auto_class.register(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)
                                  # Trying to register something existing in the Transformers library will raise an error
                                  with self.assertRaises(SCREAMING_SNAKE_CASE	):
                                         auto_class.register(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                                  # Now that the config is registered, it can be used as any other config with the auto-API
                                  _a    : str     					=BertModelTester(self	).get_config()
                                  _a    : Optional[Any]     					=NewModelConfig(**tiny_config.to_dict()	)
                                  _a    : Union[str, Any]     					=auto_class.from_config(SCREAMING_SNAKE_CASE	)
                                  self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

                                  with tempfile.TemporaryDirectory() as tmp_dir:
                                         model.save_pretrained(SCREAMING_SNAKE_CASE	)
                                         _a    : str     					=auto_class.from_pretrained(SCREAMING_SNAKE_CASE	)
                                         self.assertIsInstance(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)

             finally:
                    if "new-model" in CONFIG_MAPPING._extra_content:
                           del CONFIG_MAPPING._extra_content["new-model"]
                    for mapping in (
                        TF_MODEL_MAPPING,
                        TF_MODEL_FOR_PRETRAINING_MAPPING,
                        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
                        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
                        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
                        TF_MODEL_FOR_CAUSAL_LM_MAPPING,
                        TF_MODEL_FOR_MASKED_LM_MAPPING,
                    ):
                           if NewModelConfig in mapping._extra_content:
                                  del mapping._extra_content[NewModelConfig]






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




             '''simple docstring'''
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE     ,						"""bert-base is not a local folder and is not a valid model identifier"""	):
                    _a    : List[str]     					=TFAutoModel.from_pretrained("""bert-base"""	)






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					str:




             '''simple docstring'''
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE     ,						r"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""	):
                    _a    : Union[str, Any]     					=TFAutoModel.from_pretrained(SCREAMING_SNAKE_CASE     ,						revision="""aaaaaa"""	)






      def __UpperCAmelCase       ( self			:str	)       ->					Optional[int]:




             '''simple docstring'''
             with self.assertRaisesRegex(
                 SCREAMING_SNAKE_CASE     ,						"""hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin"""     ,						):
                    _a    : List[str]     					=TFAutoModel.from_pretrained("""hf-internal-testing/config-no-model"""	)






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




             '''simple docstring'''
             with self.assertRaisesRegex(SCREAMING_SNAKE_CASE     ,						"""Use `from_pt=True` to load this model"""	):
                    _a    : str     					=TFAutoModel.from_pretrained("""hf-internal-testing/tiny-bert-pt-only"""	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					Optional[Any]:




             '''simple docstring'''
             # Make sure we have cached the model.
             _a    : List[Any]     					=TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""	)
             with RequestCounter() as counter:
                    _a    : str     					=TFAutoModel.from_pretrained("""hf-internal-testing/tiny-random-bert"""	)
                    self.assertEqual(counter.get_request_count     ,						0	)
                    self.assertEqual(counter.head_request_count     ,						1	)
                    self.assertEqual(counter.other_request_count     ,						0	)

             # With a sharded checkpoint
             _a    : Optional[int]     					=TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""	)
             with RequestCounter() as counter:
                    _a    : List[Any]     					=TFAutoModel.from_pretrained("""ArthurZ/tiny-random-bert-sharded"""	)
                    self.assertEqual(counter.get_request_count     ,						0	)
                    self.assertEqual(counter.head_request_count     ,						1	)
                    self.assertEqual(counter.other_request_count     ,						0	)






'''simple docstring'''







import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs


@pytest.mark.parametrize(
    """kwargs, expected"""    ,[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(_UpperCAmelCase    ,i + 1  ) for i in range(10  )]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1  )]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0    ,4  ), range(4    ,7  ), range(7    ,10  )]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0    ,1  ), range(1    ,2  ), range(2    ,3  )]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Union[str, Any]    ,_UpperCAmelCase	:	Dict  )  ->			Optional[Any]:
       _a    : Tuple     					=_distribute_shards(**_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected"""    ,[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Tuple    ,_UpperCAmelCase	:	List[Any]    ,_UpperCAmelCase	:	Union[str, Any]  )  ->			List[str]:
       _a    : List[str]     					=_split_gen_kwargs(_UpperCAmelCase    ,_UpperCAmelCase  )
       assert out == expected


@pytest.mark.parametrize(
    """gen_kwargs, expected"""    ,[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ]    ,)
def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	List[Any]  )  ->			Union[str, Any]:
       if expected is RuntimeError:
              with pytest.raises(_UpperCAmelCase  ):
                     _number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
       else:
              _a    : Dict     					=_number_of_shards_in_gen_kwargs(_UpperCAmelCase  )
              assert out == expected


