





from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCamelCase__   : str					  =							logging.get_logger(__name__)


if is_torch_available():
				import torch
				from torch import nn
				from torch.utils.data import Dataset

				from transformers import Trainer







				class   _UpperCamelCase  (       a__  ):







									'''simple docstring'''

									def __init__(						self    :  List[Any] ,					__lowercase    :  Optional[int] = 1_01		):






										'''simple docstring'''







										UpperCAmelCase_								=     length





									def __len__(						self    :  str		):






										'''simple docstring'''







										return self.length





									def __getitem__(						self    :  List[str] ,					__lowercase    :  Any		):






										'''simple docstring'''







										return i







				class   _UpperCamelCase  :







									'''simple docstring'''

									def __call__(						self    :  Any ,					__lowercase    :  str		):






										'''simple docstring'''







										return {"input_ids": torch.tensor(lowerCAmelCase__		), "labels": torch.tensor(lowerCAmelCase__		)}







				class   _UpperCamelCase  (       nn.Module  ):







									'''simple docstring'''

									def __init__(						self    :  Union[str, Any]		):






										'''simple docstring'''







										super().__init__()
										# Add some (unused) params otherwise DDP will complain.
										UpperCAmelCase_								=     nn.Linear(1_20 ,					80		)





									def        SCREAMING_SNAKE_CASE							(						self    :  List[Any] ,					__lowercase    :  Dict ,					__lowercase    :  List[str]=None		):






										'''simple docstring'''







										if labels is not None:
											return torch.tensor(0.0 ,					device=input_ids.device		), input_ids
										else:
											return input_ids







class   _UpperCamelCase  (       a__  ):







					'''simple docstring'''

					@require_torch_neuroncore
					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple		):






						'''simple docstring'''







						UpperCAmelCase_								=     F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
						UpperCAmelCase_								=     self.get_auto_remove_tmp_dir()
						UpperCAmelCase_								=     F"""--output_dir {output_dir}""".split()
						UpperCAmelCase_								=     ["torchrun"] + distributed_args + args
						execute_subprocess_async(lowerCAmelCase__ ,					env=self.get_env()		)
						# successful return here == success - any errors would have caused an error in the sub-call







class   _UpperCamelCase  (       a__  ):







					'''simple docstring'''

					@require_torch_multi_gpu
					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						UpperCAmelCase_								=     F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
						UpperCAmelCase_								=     self.get_auto_remove_tmp_dir()
						UpperCAmelCase_								=     F"""--output_dir {output_dir}""".split()
						UpperCAmelCase_								=     ["torchrun"] + distributed_args + args
						execute_subprocess_async(lowerCAmelCase__ ,					env=self.get_env()		)
						# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
				# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
				#
				# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

				UpperCamelCase__   : List[str]					  =							HfArgumentParser((TrainingArguments,))
				UpperCamelCase__   : Tuple					  =							parser.parse_args_into_dataclasses()[0]

				logger.warning(
				    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, "
				    f"distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"
				)

				# Essentially, what we want to verify in the distributed case is that we get all samples back,
				# in the right order. (this is crucial for prediction for instance)
				for dataset_length in [1_01, 40, 7]:
								UpperCamelCase__   : List[Any]					  =							DummyDataset(dataset_length)




								def    A_(			A   ):
									UpperCAmelCase_								=     list(range(len(_lowercase   )   )   )
									UpperCAmelCase_								=     p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
									if not success and training_args.local_rank == 0:
										logger.warning(
										    """Predictions and/or labels do not match expected results:\n  - predictions: """
										    f"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""   )
									return {"success": success}

								UpperCamelCase__   : List[str]					  =							Trainer(
								    model=DummyModel(),
								    args=training_args,
								    data_collator=DummyDataCollator(),
								    eval_dataset=dataset,
								    compute_metrics=compute_metrics,
								)
								UpperCamelCase__   : Tuple					  =							trainer.evaluate()
								logger.info(metrics)
								if metrics["eval_success"] is not True:
												logger.error(metrics)
												exit(1)

								UpperCamelCase__   : int					  =							trainer.predict(dataset)
								logger.info(p.metrics)
								if p.metrics["test_success"] is not True:
												logger.error(p.metrics)
												exit(1)

								UpperCamelCase__   : Optional[Any]					  =							2

								UpperCamelCase__   : Optional[int]					  =							trainer.evaluate()
								logger.info(metrics)
								if metrics["eval_success"] is not True:
												logger.error(metrics)
												exit(1)

								UpperCamelCase__   : Optional[int]					  =							trainer.predict(dataset)
								logger.info(p.metrics)
								if p.metrics["test_success"] is not True:
												logger.error(p.metrics)
												exit(1)
								UpperCamelCase__   : List[Any]					  =							None









from __future__ import annotations

import unittest

from transformers import is_tf_available, is_torch_available
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, is_pt_tf_cross_test, slow


if is_tf_available():
				from transformers import (
				    AutoConfig,
				    BertConfig,
				    GPTaConfig,
				    TaConfig,
				    TFAutoModel,
				    TFAutoModelForCausalLM,
				    TFAutoModelForMaskedLM,
				    TFAutoModelForPreTraining,
				    TFAutoModelForQuestionAnswering,
				    TFAutoModelForSeqaSeqLM,
				    TFAutoModelForSequenceClassification,
				    TFAutoModelWithLMHead,
				    TFBertForMaskedLM,
				    TFBertForPreTraining,
				    TFBertForQuestionAnswering,
				    TFBertForSequenceClassification,
				    TFBertModel,
				    TFGPTaLMHeadModel,
				    TFRobertaForMaskedLM,
				    TFTaForConditionalGeneration,
				)
				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

if is_torch_available():
				from transformers import (
				    AutoModel,
				    AutoModelForCausalLM,
				    AutoModelForMaskedLM,
				    AutoModelForPreTraining,
				    AutoModelForQuestionAnswering,
				    AutoModelForSeqaSeqLM,
				    AutoModelForSequenceClassification,
				    AutoModelWithLMHead,
				    BertForMaskedLM,
				    BertForPreTraining,
				    BertForQuestionAnswering,
				    BertForSequenceClassification,
				    BertModel,
				    GPTaLMHeadModel,
				    RobertaForMaskedLM,
				    TaForConditionalGeneration,
				)







@is_pt_tf_cross_test
class   _UpperCamelCase  (       unittest.TestCase  ):







					'''simple docstring'''

					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						for model_name in ["bert-base-uncased"]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModel.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModel.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						for model_name in ["bert-base-uncased"]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForPreTraining.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForPreTraining.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Any		):






						'''simple docstring'''







						for model_name in TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForCausalLM.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     TFAutoModelForCausalLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForCausalLM.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     AutoModelForCausalLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelWithLMHead.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelWithLMHead.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						for model_name in TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForMaskedLM.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     TFAutoModelForMaskedLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForMaskedLM.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     AutoModelForMaskedLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Union[str, Any]		):






						'''simple docstring'''







						for model_name in TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForSeqaSeqLM.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     TFAutoModelForSeqaSeqLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							UpperCAmelCase_						,	UpperCAmelCase_								=     AutoModelForSeqaSeqLM.from_pretrained(
							    __lowercase ,					output_loading_info=__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[int]		):






						'''simple docstring'''







						for model_name in ["bert-base-uncased"]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForSequenceClassification.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForSequenceClassification.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					@slow
					def        SCREAMING_SNAKE_CASE							(						self    :  str		):






						'''simple docstring'''







						for model_name in ["bert-base-uncased"]:
							UpperCAmelCase_								=     AutoConfig.from_pretrained(__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     TFAutoModelForQuestionAnswering.from_pretrained(__lowercase ,					from_pt=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)

							UpperCAmelCase_								=     AutoModelForQuestionAnswering.from_pretrained(__lowercase ,					from_tf=__lowercase		)
							self.assertIsNotNone(__lowercase		)
							self.assertIsInstance(__lowercase ,					__lowercase		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Tuple		):






						'''simple docstring'''







						UpperCAmelCase_								=     TFAutoModelWithLMHead.from_pretrained(__lowercase ,					from_pt=__lowercase		)
						self.assertIsInstance(__lowercase ,					__lowercase		)
						self.assertEqual(model.num_parameters() ,					1_44_10		)
						self.assertEqual(model.num_parameters(only_trainable=__lowercase		) ,					1_44_10		)

						UpperCAmelCase_								=     AutoModelWithLMHead.from_pretrained(__lowercase ,					from_tf=__lowercase		)
						self.assertIsInstance(__lowercase ,					__lowercase		)
						self.assertEqual(model.num_parameters() ,					1_44_10		)
						self.assertEqual(model.num_parameters(only_trainable=__lowercase		) ,					1_44_10		)





					def        SCREAMING_SNAKE_CASE							(						self    :  Optional[Any]		):






						'''simple docstring'''







						UpperCAmelCase_								=     TFAutoModelWithLMHead.from_pretrained(__lowercase ,					from_pt=__lowercase		)
						self.assertIsInstance(__lowercase ,					__lowercase		)
						self.assertEqual(model.num_parameters() ,					1_44_10		)
						self.assertEqual(model.num_parameters(only_trainable=__lowercase		) ,					1_44_10		)

						UpperCAmelCase_								=     AutoModelWithLMHead.from_pretrained(__lowercase ,					from_tf=__lowercase		)
						self.assertIsInstance(__lowercase ,					__lowercase		)
						self.assertEqual(model.num_parameters() ,					1_44_10		)
						self.assertEqual(model.num_parameters(only_trainable=__lowercase		) ,					1_44_10		)


