





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


A_												=     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 SCREAMING_SNAKE_CASE_ (					lowercase_   ):





									"""simple docstring"""


									def __init__( self	,  _lowerCAmelCase = 101	):
														lowerCamelCase__       =       length





									def __len__( self	):
														return self.length






									def __getitem__( self	,  _lowerCAmelCase	):
														return i

				class SCREAMING_SNAKE_CASE_ :





									"""simple docstring"""


									def __call__( self	,  _lowerCAmelCase	):
														return {"input_ids": torch.tensor(_lowerCAmelCase	), "labels": torch.tensor(_lowerCAmelCase	)}

				class SCREAMING_SNAKE_CASE_ (					nn.Module   ):





									"""simple docstring"""


									def __init__( self	):
														super().__init__()
														# Add some (unused) params otherwise DDP will complain.
														lowerCamelCase__       =       nn.Linear(120	,  80	)






									def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase=None	):







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

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@require_torch_neuroncore
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       F"--nproc_per_node=2\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
										lowerCamelCase__       =       self.get_auto_remove_tmp_dir()
										lowerCamelCase__       =       F"--output_dir {output_dir}".split()
										lowerCamelCase__       =       ['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 SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@require_torch_multi_gpu
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       F"--nproc_per_node={torch.cuda.device_count()}\n            --master_port={get_torch_dist_unique_port()}\n            {self.test_file_dir}/test_trainer_distributed.py\n        ".split()
										lowerCamelCase__       =       self.get_auto_remove_tmp_dir()
										lowerCamelCase__       =       F"--output_dir {output_dir}".split()
										lowerCamelCase__       =       ['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

				A_												=     HfArgumentParser((TrainingArguments,))
				A_												=     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_0_1, 4_0, 7]:
								A_												=     DummyDataset(dataset_length)





								def    __UpperCamelCase      (    a) ->Dict:
													lowerCamelCase__       =       list(range(len(SCREAMING_SNAKE_CASE_)))
													lowerCamelCase__       =       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}

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

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

								A_												=     2

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

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







								A_												=     None








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig


A_												=     {
    "albert-base-v1": "https://huggingface.co/albert-base-v1/resolve/main/config.json",
    "albert-large-v1": "https://huggingface.co/albert-large-v1/resolve/main/config.json",
    "albert-xlarge-v1": "https://huggingface.co/albert-xlarge-v1/resolve/main/config.json",
    "albert-xxlarge-v1": "https://huggingface.co/albert-xxlarge-v1/resolve/main/config.json",
    "albert-base-v2": "https://huggingface.co/albert-base-v2/resolve/main/config.json",
    "albert-large-v2": "https://huggingface.co/albert-large-v2/resolve/main/config.json",
    "albert-xlarge-v2": "https://huggingface.co/albert-xlarge-v2/resolve/main/config.json",
    "albert-xxlarge-v2": "https://huggingface.co/albert-xxlarge-v2/resolve/main/config.json",
}

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					A__				 =    "albert"





					def __init__( self	,  _lowerCAmelCase=3_0000	,  _lowerCAmelCase=128	,  _lowerCAmelCase=4096	,  _lowerCAmelCase=12	,  _lowerCAmelCase=1	,  _lowerCAmelCase=64	,  _lowerCAmelCase=1_6384	,  _lowerCAmelCase=1	,  _lowerCAmelCase="gelu_new"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=0	,  _lowerCAmelCase=512	,  _lowerCAmelCase=2	,  _lowerCAmelCase=0.02	,  _lowerCAmelCase=1E-12	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase="absolute"	,  _lowerCAmelCase=0	,  _lowerCAmelCase=2	,  _lowerCAmelCase=3	,  **_lowerCAmelCase	,  ):
										super().__init__(pad_token_id=_lowerCAmelCase	,  bos_token_id=_lowerCAmelCase	,  eos_token_id=_lowerCAmelCase	,  **_lowerCAmelCase	)

										lowerCamelCase__       =       vocab_size
										lowerCamelCase__       =       embedding_size
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_hidden_groups
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       inner_group_num
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       max_position_embeddings
										lowerCamelCase__       =       type_vocab_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       layer_norm_eps
										lowerCamelCase__       =       classifier_dropout_prob
										lowerCamelCase__       =       position_embedding_type

class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@property
					def 					__magic_name__ ( self	):
										if self.task == "multiple-choice":
															lowerCamelCase__       =       {0: "batch", 1: "choice", 2: "sequence"}
										else:
															lowerCamelCase__       =       {0: "batch", 1: "sequence"}
										return OrderedDict(
										    [
										        ("input_ids", dynamic_axis),
										        ("attention_mask", dynamic_axis),
										        ("token_type_ids", dynamic_axis),
										    ]	)

