





import inspect
import os
import unittest

import torch

import accelerate
from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_multi_gpu
from accelerate.utils import patch_environment
class SCREAMING_SNAKE_CASE_						(					unittest.TestCase					):






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







							"""simple docstring"""


							snake_case_		:			Optional[Any]											=	inspect.getfile(accelerate.test_utils	)
							snake_case_		:			Any											=	os.path.sep.join(mod_file.split(os.path.sep	)[:-1] + ['scripts', 'test_script.py']	)
							snake_case_		:			Optional[int]											=	os.path.sep.join(
							    mod_file.split(os.path.sep	)[:-1] + ['scripts', 'test_distributed_data_loop.py']	)
							snake_case_		:			Tuple											=	os.path.sep.join(mod_file.split(os.path.sep	)[:-1] + ['scripts', 'test_ops.py']	)






						@require_multi_gpu
						def UpperCAmelCase_				(      self      :      List[str]	)							->    Any:







							"""simple docstring"""


							print(F"""Found {torch.cuda.device_count()} devices."""	)
							snake_case_		:			str											=	['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.test_file_path]
							with patch_environment(omp_num_threads=1	):
								execute_subprocess_async(_A		,   env=os.environ.copy()	)






						@require_multi_gpu
						def UpperCAmelCase_				(      self      :      Tuple	)							->    str:







							"""simple docstring"""


							print(F"""Found {torch.cuda.device_count()} devices."""	)
							snake_case_		:			Optional[int]											=	['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.operation_file_path]
							print(F"""Command: {cmd}"""	)
							with patch_environment(omp_num_threads=1	):
								execute_subprocess_async(_A		,   env=os.environ.copy()	)






						@require_multi_gpu
						def UpperCAmelCase_				(      self      :      Optional[int]	)							->    Optional[int]:







							"""simple docstring"""


							snake_case_		:			Optional[int]											=	['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__	)]
							with patch_environment(omp_num_threads=1	):
								execute_subprocess_async(_A		,   env=os.environ.copy()	)






						@require_multi_gpu
						def UpperCAmelCase_				(      self      :      Dict	)							->    Optional[Any]:







							"""simple docstring"""


							print(F"""Found {torch.cuda.device_count()} devices, using 2 devices only"""	)
							snake_case_		:			Dict											=	['torchrun', F"""--nproc_per_node={torch.cuda.device_count()}""", self.data_loop_file_path]
							with patch_environment(omp_num_threads=1		,   cuda_visible_devices='0,1'	):
								execute_subprocess_async(_A		,   env=os.environ.copy()	)


if __name__ == "__main__":
						_SCREAMING_SNAKE_CASE      	=		Accelerator()
						_SCREAMING_SNAKE_CASE      	=		(accelerator.state.process_index + 2, 10)
						_SCREAMING_SNAKE_CASE      	=		torch.randint(0, 10, shape).to(accelerator.device)

						_SCREAMING_SNAKE_CASE      	=		""""""

						_SCREAMING_SNAKE_CASE      	=		accelerator.pad_across_processes(tensor)
						if tensora.shape[0] != accelerator.state.num_processes + 1:
												error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
						if not torch.equal(tensora[: accelerator.state.process_index + 2], tensor):
												error_msg += "Tensors have different values."
						if not torch.all(tensora[accelerator.state.process_index + 2 :] == 0):
												error_msg += "Padding was not done with the right value (0)."

						_SCREAMING_SNAKE_CASE      	=		accelerator.pad_across_processes(tensor, pad_first=True)
						if tensora.shape[0] != accelerator.state.num_processes + 1:
												error_msg += F"Found shape {tensora.shape} but should have {accelerator.state.num_processes + 1} at dim 0."
						_SCREAMING_SNAKE_CASE      	=		accelerator.state.num_processes - accelerator.state.process_index - 1
						if not torch.equal(tensora[index:], tensor):
												error_msg += "Tensors have different values."
						if not torch.all(tensora[:index] == 0):
												error_msg += "Padding was not done with the right value (0)."

						# Raise error at the end to make sure we don't stop at the first failure.
						if len(error_msg) > 0:
												raise ValueError(error_msg)








import argparse
import json

import numpy
import torch

from transformers.models.xlm.tokenization_xlm import VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()


def 							SCREAMING_SNAKE_CASE__						(    __a   , __a   ):
	# Load checkpoint
	snake_case_		:			Union[str, Any]											=	torch.load(__a   , map_location='cpu'   )

	snake_case_		:			Union[str, Any]											=	chkpt['model']

	# We have the base model one level deeper than the original XLM repository
	snake_case_		:			str											=	{}
	for k, v in state_dict.items():
		if "pred_layer" in k:
			snake_case_		:			Tuple											=	v
		else:
			snake_case_		:			Dict											=	v

	snake_case_		:			Tuple											=	chkpt['params']
	snake_case_		:			List[Any]											=	{n: v for n, v in config.items() if not isinstance(__a   , (torch.FloatTensor, numpy.ndarray)   )}

	snake_case_		:			Optional[int]											=	chkpt['dico_word2id']
	snake_case_		:			List[str]											=	{s + '</w>' if s.find('@@'   ) == -1 and i > 13 else s.replace('@@'   , ''   ): i for s, i in vocab.items()}

	# Save pytorch-model
	snake_case_		:			List[str]											=	pytorch_dump_folder_path + '/' + WEIGHTS_NAME
	snake_case_		:			Dict											=	pytorch_dump_folder_path + '/' + CONFIG_NAME
	snake_case_		:			Any											=	pytorch_dump_folder_path + '/' + VOCAB_FILES_NAMES['vocab_file']

	print(f"""Save PyTorch model to {pytorch_weights_dump_path}"""   )
	torch.save(__a   , __a   )

	print(f"""Save configuration file to {pytorch_config_dump_path}"""   )
	with open(__a   , 'w'   , encoding='utf-8'   ) as f:
		f.write(json.dumps(__a   , indent=2   ) + '\n'   )

	print(f"""Save vocab file to {pytorch_config_dump_path}"""   )
	with open(__a   , 'w'   , encoding='utf-8'   ) as f:
		f.write(json.dumps(__a   , indent=2   ) + '\n'   )


if __name__ == "__main__":
						_SCREAMING_SNAKE_CASE      	=		argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--xlm_checkpoint_path""", default=None, type=str, required=True, help="""Path the official PyTorch dump."""
						)
						parser.add_argument(
						    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
						)
						_SCREAMING_SNAKE_CASE      	=		parser.parse_args()
						convert_xlm_checkpoint_to_pytorch(args.xlm_checkpoint_path, args.pytorch_dump_folder_path)

