

"""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 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					lowercase__												=			"new-model"


if is_tf_available():







							class 							SCREAMING_SNAKE_CASE   (						a_			):






												"""simple docstring"""



												lowercase__												=			NewModelConfig







@require_tf
class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					@slow
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						lowerCAmelCase__      :		str    							=      '''bert-base-cased'''
						lowerCAmelCase__      :		Any    							=      AutoConfig.from_pretrained(lowercase_						)
						self.assertIsNotNone(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)

						lowerCAmelCase__      :		Optional[Any]    							=      TFAutoModel.from_pretrained(lowercase_						)
						self.assertIsNotNone(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						lowerCAmelCase__      :		List[Any]    							=      '''bert-base-cased'''
						lowerCAmelCase__      :		Any    							=      AutoConfig.from_pretrained(lowercase_						)
						self.assertIsNotNone(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)

						lowerCAmelCase__      :		Union[str, Any]    							=      TFAutoModelForPreTraining.from_pretrained(lowercase_						)
						self.assertIsNotNone(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							str						):
						for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCAmelCase__      :		Dict    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Union[str, Any]    							=      TFAutoModelForCausalLM.from_pretrained(lowercase_						)
							lowerCAmelCase__							,lowerCAmelCase__      :		List[str]    							=      TFAutoModelForCausalLM.from_pretrained(lowercase_  ,output_loading_info=lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCAmelCase__      :		Dict    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		int    							=      TFAutoModelWithLMHead.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Any						):
						for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCAmelCase__      :		List[str]    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Union[str, Any]    							=      TFAutoModelForMaskedLM.from_pretrained(lowercase_						)
							lowerCAmelCase__							,lowerCAmelCase__      :		Optional[int]    							=      TFAutoModelForMaskedLM.from_pretrained(lowercase_  ,output_loading_info=lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Any						):
						for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							lowerCAmelCase__      :		List[str]    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Any    							=      TFAutoModelForSeqaSeqLM.from_pretrained(lowercase_						)
							lowerCAmelCase__							,lowerCAmelCase__      :		Optional[Any]    							=      TFAutoModelForSeqaSeqLM.from_pretrained(lowercase_  ,output_loading_info=lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Tuple						):
						# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						for model_name in ["bert-base-uncased"]:
							lowerCAmelCase__      :		Optional[Any]    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Dict    							=      TFAutoModelForSequenceClassification.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					def 			__lowerCAmelCase					(			self     :							Any						):
						# for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						for model_name in ["bert-base-uncased"]:
							lowerCAmelCase__      :		Dict    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Union[str, Any]    							=      TFAutoModelForQuestionAnswering.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					@slow
					@require_tensorflow_probability
					def 			__lowerCAmelCase					(			self     :							Dict						):
						for model_name in TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST[5:6]:
							lowerCAmelCase__      :		List[Any]    							=      AutoConfig.from_pretrained(lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)

							lowerCAmelCase__      :		Any    							=      TFAutoModelForTableQuestionAnswering.from_pretrained(lowercase_						)
							lowerCAmelCase__							,lowerCAmelCase__      :		Dict    							=      TFAutoModelForTableQuestionAnswering.from_pretrained(
							    lowercase_  ,output_loading_info=lowercase_						)
							self.assertIsNotNone(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						lowerCAmelCase__      :		str    							=      TFAutoModelWithLMHead.from_pretrained(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)
						self.assertEqual(model.num_parameters()  ,1_4_4_1_0						)
						self.assertEqual(model.num_parameters(only_trainable=lowercase_						)  ,1_4_4_1_0						)







					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						lowerCAmelCase__      :		List[Any]    							=      TFAutoModelWithLMHead.from_pretrained(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)
						self.assertEqual(model.num_parameters()  ,1_4_4_1_0						)
						self.assertEqual(model.num_parameters(only_trainable=lowercase_						)  ,1_4_4_1_0						)







					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						# For the auto model mapping, FunnelConfig has two models: FunnelModel and FunnelBaseModel
						lowerCAmelCase__      :		List[Any]    							=      TFAutoModel.from_pretrained('''sgugger/funnel-random-tiny'''						)
						self.assertIsInstance(lowercase_  ,lowercase_						)

						lowerCAmelCase__      :		Tuple    							=      copy.deepcopy(model.config						)
						lowerCAmelCase__      :		Optional[Any]    							=      ['''FunnelBaseModel''']
						lowerCAmelCase__      :		Dict    							=      TFAutoModel.from_config(lowercase_						)
						self.assertIsInstance(lowercase_  ,lowercase_						)

						with tempfile.TemporaryDirectory() as tmp_dir:
							model.save_pretrained(lowercase_						)
							lowerCAmelCase__      :		str    							=      TFAutoModel.from_pretrained(lowercase_						)
							self.assertIsInstance(lowercase_  ,lowercase_						)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						try:
							AutoConfig.register('''new-model'''  ,lowercase_						)

							lowerCAmelCase__      :		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(lowercase_						):
										auto_class.register(lowercase_  ,lowercase_						)
									auto_class.register(lowercase_  ,lowercase_						)
									# Trying to register something existing in the Transformers library will raise an error
									with self.assertRaises(lowercase_						):
										auto_class.register(lowercase_  ,lowercase_						)

									# Now that the config is registered, it can be used as any other config with the auto-API
									lowerCAmelCase__      :		str    							=      BertModelTester(self						).get_config()
									lowerCAmelCase__      :		int    							=      NewModelConfig(**tiny_config.to_dict()						)
									lowerCAmelCase__      :		Tuple    							=      auto_class.from_config(lowercase_						)
									self.assertIsInstance(lowercase_  ,lowercase_						)

									with tempfile.TemporaryDirectory() as tmp_dir:
										model.save_pretrained(lowercase_						)
										lowerCAmelCase__      :		Optional[Any]    							=      auto_class.from_pretrained(lowercase_						)
										self.assertIsInstance(lowercase_  ,lowercase_						)

						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 			__lowerCAmelCase					(			self     :							Dict						):
						with self.assertRaisesRegex(
						    lowercase_  ,'''bert-base is not a local folder and is not a valid model identifier'''						):
							lowerCAmelCase__      :		int    							=      TFAutoModel.from_pretrained('''bert-base'''						)







					def 			__lowerCAmelCase					(			self     :							List[str]						):
						with self.assertRaisesRegex(
						    lowercase_  ,R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''						):
							lowerCAmelCase__      :		Dict    							=      TFAutoModel.from_pretrained(lowercase_  ,revision='''aaaaaa'''						)







					def 			__lowerCAmelCase					(			self     :							Optional[Any]						):
						with self.assertRaisesRegex(
						    lowercase_  ,'''hf-internal-testing/config-no-model does not appear to have a file named pytorch_model.bin'''  ,):
							lowerCAmelCase__      :		Dict    							=      TFAutoModel.from_pretrained('''hf-internal-testing/config-no-model'''						)







					def 			__lowerCAmelCase					(			self     :							str						):
						with self.assertRaisesRegex(lowercase_  ,'''Use `from_pt=True` to load this model'''						):
							lowerCAmelCase__      :		Optional[Any]    							=      TFAutoModel.from_pretrained('''hf-internal-testing/tiny-bert-pt-only'''						)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						# Make sure we have cached the model.
						lowerCAmelCase__      :		int    							=      TFAutoModel.from_pretrained('''hf-internal-testing/tiny-random-bert'''						)
						with RequestCounter() as counter:
							lowerCAmelCase__      :		int    							=      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
						lowerCAmelCase__      :		Union[str, Any]    							=      TFAutoModel.from_pretrained('''ArthurZ/tiny-random-bert-sharded'''						)
						with RequestCounter() as counter:
							lowerCAmelCase__      :		Optional[int]    							=      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 os
from argparse import ArgumentParser
from typing import List

import torch.utils.data

from datasets import Dataset, IterableDataset
from datasets.distributed import split_dataset_by_node


__UpperCamelCase    : str				 =						4
__UpperCamelCase    : List[str]				 =						3







class 							SCREAMING_SNAKE_CASE   (						a_			):






					"""simple docstring"""



					pass

def     __SCREAMING_SNAKE_CASE    ( A_					):
	for shard in shards:
		for i in range(A_					):
			yield {"i": i, "shard": shard}

def     __SCREAMING_SNAKE_CASE    ( ):
	lowerCAmelCase__      :		List[str]    							=      int(os.environ['''RANK''']					)
	lowerCAmelCase__      :		Union[str, Any]    							=      int(os.environ['''WORLD_SIZE''']					)

	lowerCAmelCase__      :		Tuple    							=      ArgumentParser()
	parser.add_argument('''--streaming'''		,	type=A_					)
	parser.add_argument('''--local_rank'''		,	type=A_					)
	parser.add_argument('''--num_workers'''		,	type=A_		,	default=0					)
	lowerCAmelCase__      :		Tuple    							=      parser.parse_args()
	lowerCAmelCase__      :		Tuple    							=      args.streaming
	lowerCAmelCase__      :		Any    							=      args.num_workers

	lowerCAmelCase__      :		Dict    							=      {'''shards''': [f'shard_{shard_idx}' for shard_idx in range(A_					)]}
	lowerCAmelCase__      :		Optional[int]    							=      IterableDataset.from_generator(A_		,	gen_kwargs=A_					)
	if not streaming:
		lowerCAmelCase__      :		Dict    							=      Dataset.from_list(list(A_					)					)

	lowerCAmelCase__      :		Optional[Any]    							=      split_dataset_by_node(A_		,	rank=A_		,	world_size=A_					)
	lowerCAmelCase__      :		List[Any]    							=      torch.utils.data.DataLoader(A_		,	num_workers=A_					)

	lowerCAmelCase__      :		Optional[int]    							=      NUM_SHARDS * NUM_ITEMS_PER_SHARD
	lowerCAmelCase__      :		List[str]    							=      full_size // world_size
	expected_local_size += int(rank < (full_size % world_size)					)

	lowerCAmelCase__      :		int    							=      sum(1 for _ in dataloader					)
	if local_size != expected_local_size:
		raise FailedTestError(f'local_size {local_size} != expected_local_size {expected_local_size}'					)


if __name__ == "__main__":
							main()




