






'''simple docstring'''
import inspect
import logging
import os
import random
import shutil
import tempfile
import unittest

import pytest
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

from accelerate import Accelerator
from accelerate.test_utils import execute_subprocess_async, require_cuda
from accelerate.utils import ProjectConfiguration, set_seed


UpperCamelCase__ =			logging.getLogger(__name__)




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase=2						,		_UpperCamelCase=3						,		_UpperCamelCase=16						,		_UpperCamelCase = 10						,		_UpperCamelCase = 2							):


			"""simple docstring"""




			def get_dataset(_UpperCamelCase							):
						lowercase_	:    Tuple           =     torch.randn(batch_size * n_batches						,		1							)
						return TensorDataset(__lowerCAmelCase						,		a * x + b + 0.1 * torch.randn(batch_size * n_batches						,		1							)							)

			lowercase_	:    Dict           =     get_dataset(__lowerCAmelCase							)
			lowercase_	:    Optional[int]           =     get_dataset(__lowerCAmelCase							)
			lowercase_	:    Any           =     DataLoader(__lowerCAmelCase						,		shuffle=__lowerCAmelCase						,		batch_size=__lowerCAmelCase						,		num_workers=4							)
			lowercase_	:    str           =     DataLoader(__lowerCAmelCase						,		shuffle=__lowerCAmelCase						,		batch_size=__lowerCAmelCase						,		num_workers=4							)
			return (train_dataloader, valid_dataloader)




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase=None							):


			"""simple docstring"""




			lowercase_	:    Optional[int]           =     []
			for epoch in range(__lowerCAmelCase							):
						# Train quickly
						model.train()
						for batch in dataloader:
									lowercase_	:    Optional[int]           =     batch
									lowercase_	:    Optional[int]           =     model(__lowerCAmelCase							)
									lowercase_	:    Optional[int]           =     torch.nn.functional.mse_loss(__lowerCAmelCase						,		__lowerCAmelCase							)
									accelerator.backward(__lowerCAmelCase							)
									optimizer.step()
									optimizer.zero_grad()
						rands.append(random.random()							)  # Introduce some randomness
						if scheduler is not None:
									scheduler.step()
			return rands


class        _UpperCAmelCase						(		nn.Module			):




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






							'''simple docstring'''






							super().__init__()
							lowercase_	:    Any           =     nn.Parameter(torch.randn(1			)			)
							lowercase_	:    Tuple           =     nn.Parameter(torch.randn(1			)			)





				def 				lowerCAmelCase__				(			self  :       str  ,     a  :       Optional[int]			):






							'''simple docstring'''






							return x * self.a + self.b


class        _UpperCAmelCase						(		unittest.TestCase			):




				def 				lowerCAmelCase__				(			self  :       Any			):






							'''simple docstring'''






							with tempfile.TemporaryDirectory() as tmpdir:
										set_seed(4_2			)
										lowercase_	:    List[Any]           =     DummyModel()
										lowercase_	:    List[str]           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    Union[str, Any]           =     dummy_dataloaders()
										lowercase_	:    Tuple           =     ProjectConfiguration(total_limit=1  ,     project_dir=lowerCamelCase__  ,     automatic_checkpoint_naming=lowerCamelCase__			)
										# Train baseline
										lowercase_	:    List[str]           =     Accelerator(project_config=lowerCamelCase__			)
										lowercase_	:    Optional[Any]           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save initial
										accelerator.save_state()

										# Save second state
										accelerator.save_state()
										self.assertEqual(len(os.listdir(accelerator.project_dir			)			)  ,     1			)




				def 				lowerCAmelCase__				(			self  :       Any			):






							'''simple docstring'''






							with tempfile.TemporaryDirectory() as tmpdir:
										set_seed(4_2			)
										lowercase_	:    List[Any]           =     DummyModel()
										lowercase_	:    Dict           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    Union[str, Any]           =     dummy_dataloaders()
										# Train baseline
										lowercase_	:    List[str]           =     Accelerator()
										lowercase_	:    Optional[Any]           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save initial
										lowercase_	:    Any           =     os.path.join(lowerCamelCase__  ,     "initial"			)
										accelerator.save_state(lowerCamelCase__			)
										(lowercase_)	:    Union[str, Any]           =     model.a.item(), model.b.item()
										lowercase_	:    str           =     optimizer.state_dict()
										lowercase_	:    Tuple           =     train(3  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										(lowercase_)	:    Optional[int]           =     model.a.item(), model.b.item()
										lowercase_	:    List[str]           =     optimizer.state_dict()

										# Train partially
										set_seed(4_2			)
										lowercase_	:    Optional[Any]           =     DummyModel()
										lowercase_	:    int           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    List[Any]           =     dummy_dataloaders()
										lowercase_	:    Union[str, Any]           =     Accelerator()
										lowercase_	:    Optional[Any]           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										accelerator.load_state(lowerCamelCase__			)
										(lowercase_)	:    Tuple           =     model.a.item(), model.b.item()
										lowercase_	:    Union[str, Any]           =     optimizer.state_dict()
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)

										lowercase_	:    Optional[Any]           =     train(2  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save everything
										lowercase_	:    Optional[Any]           =     os.path.join(lowerCamelCase__  ,     "checkpoint"			)
										accelerator.save_state(lowerCamelCase__			)

										# Load everything back in and make sure all states work
										accelerator.load_state(lowerCamelCase__			)
										test_rands += train(1  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										(lowercase_)	:    Union[str, Any]           =     model.a.item(), model.b.item()
										lowercase_	:    Dict           =     optimizer.state_dict()
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)




				def 				lowerCAmelCase__				(			self  :       Tuple			):






							'''simple docstring'''






							with tempfile.TemporaryDirectory() as tmpdir:
										set_seed(4_2			)
										lowercase_	:    Optional[Any]           =     DummyModel()
										lowercase_	:    int           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    Any           =     dummy_dataloaders()
										lowercase_	:    Union[str, Any]           =     ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__			)

										# Train baseline
										lowercase_	:    str           =     Accelerator(project_dir=lowerCamelCase__  ,     project_config=lowerCamelCase__			)
										lowercase_	:    Optional[Any]           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save initial
										accelerator.save_state()
										(lowercase_)	:    Optional[Any]           =     model.a.item(), model.b.item()
										lowercase_	:    Union[str, Any]           =     optimizer.state_dict()
										lowercase_	:    Optional[Any]           =     train(3  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										(lowercase_)	:    Dict           =     model.a.item(), model.b.item()
										lowercase_	:    Optional[Any]           =     optimizer.state_dict()

										# Train partially
										set_seed(4_2			)
										lowercase_	:    Optional[Any]           =     DummyModel()
										lowercase_	:    str           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    Dict           =     dummy_dataloaders()
										lowercase_	:    Optional[Any]           =     ProjectConfiguration(iteration=1  ,     automatic_checkpoint_naming=lowerCamelCase__			)
										lowercase_	:    Optional[Any]           =     Accelerator(project_dir=lowerCamelCase__  ,     project_config=lowerCamelCase__			)
										lowercase_	:    Union[str, Any]           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										accelerator.load_state(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_0"			)			)
										(lowercase_)	:    int           =     model.a.item(), model.b.item()
										lowercase_	:    int           =     optimizer.state_dict()
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)

										lowercase_	:    str           =     train(2  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save everything
										accelerator.save_state()

										# Load everything back in and make sure all states work
										accelerator.load_state(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_1"			)			)
										test_rands += train(1  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										(lowercase_)	:    Dict           =     model.a.item(), model.b.item()
										lowercase_	:    Dict           =     optimizer.state_dict()
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertEqual(lowerCamelCase__  ,     lowerCamelCase__			)




				def 				lowerCAmelCase__				(			self  :       Optional[Any]			):






							'''simple docstring'''






							lowercase_	:    int           =     torch.tensor([1, 2, 3]			)
							lowercase_	:    Any           =     torch.tensor([2, 3, 4]			)
							lowercase_	:    Optional[int]           =     DummyModel()
							lowercase_	:    List[Any]           =     torch.optim.Adam(net.parameters()			)
							lowercase_	:    Union[str, Any]           =     Accelerator()
							with self.assertRaises(lowerCamelCase__			) as ve:
										accelerator.register_for_checkpointing(lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
							lowercase_	:    Tuple           =     str(ve.exception			)
							self.assertTrue("Item at index 0" in message			)
							self.assertTrue("Item at index 1" in message			)
							self.assertFalse("Item at index 2" in message			)
							self.assertFalse("Item at index 3" in message			)




				def 				lowerCAmelCase__				(			self  :       List[str]			):






							'''simple docstring'''






							with tempfile.TemporaryDirectory() as tmpdir:
										set_seed(4_2			)
										lowercase_	:    Optional[int]           =     DummyModel()
										lowercase_	:    int           =     torch.optim.Adam(params=model.parameters()  ,     lr=1e-3			)
										lowercase_	:    int           =     torch.optim.lr_scheduler.StepLR(lowerCamelCase__  ,     step_size=1  ,     gamma=0.99			)
										lowercase_	:    Any           =     dummy_dataloaders()
										lowercase_	:    Any           =     ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__			)
										# Train baseline
										lowercase_	:    Optional[Any]           =     Accelerator(project_dir=lowerCamelCase__  ,     project_config=lowerCamelCase__			)
										lowercase_	:    int           =     accelerator.prepare(
										    lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										# Save initial
										accelerator.save_state()
										lowercase_	:    Union[str, Any]           =     scheduler.state_dict()
										train(3  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__  ,     lowerCamelCase__			)
										self.assertNotEqual(lowerCamelCase__  ,     scheduler.state_dict()			)

										# Load everything back in and make sure all states work
										accelerator.load_state(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_0"			)			)
										self.assertEqual(lowerCamelCase__  ,     scheduler.state_dict()			)




				def 				lowerCAmelCase__				(			self  :       List[Any]			):






							'''simple docstring'''






							with tempfile.TemporaryDirectory() as tmpdir:
										set_seed(4_2			)
										lowercase_	:    Dict           =     DummyModel()
										lowercase_	:    int           =     ProjectConfiguration(automatic_checkpoint_naming=lowerCamelCase__  ,     total_limit=2			)
										# Train baseline
										lowercase_	:    Optional[int]           =     Accelerator(project_dir=lowerCamelCase__  ,     project_config=lowerCamelCase__			)
										lowercase_	:    Union[str, Any]           =     accelerator.prepare(lowerCamelCase__			)
										# Save 3 states:
										for _ in range(1_1			):
													accelerator.save_state()
										self.assertTrue(not os.path.exists(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_0"			)			)			)
										self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_9"			)			)			)
										self.assertTrue(os.path.exists(os.path.join(lowerCamelCase__  ,     "checkpoints"  ,     "checkpoint_10"			)			)			)





				@require_cuda
				def 				lowerCAmelCase__				(			self  :       List[Any]			):






							'''simple docstring'''






							lowercase_	:    str           =     ["torchrun", f"""--nproc_per_node={torch.cuda.device_count()}""", inspect.getfile(self.__class__			)]
							execute_subprocess_async(lowerCamelCase__  ,     env=os.environ.copy()			)


if __name__ == "__main__":
	UpperCamelCase__ =			'/tmp/accelerate/state_checkpointing'
	UpperCamelCase__ =			DummyModel()
	UpperCamelCase__ =			torch.optim.Adam(params=model.parameters(), lr=1e-3)
	UpperCamelCase__ =			torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
	UpperCamelCase__,     UpperCamelCase__ =			dummy_dataloaders()
	UpperCamelCase__ =			ProjectConfiguration(automatic_checkpoint_naming=True)
	# Train baseline
	UpperCamelCase__ =			Accelerator(project_dir=savedir, project_config=project_config, mixed_precision='no')
	if accelerator.process_index == 0:
		if os.path.exists(savedir):
			shutil.rmtree(savedir)
		os.makedirs(savedir)
	UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__,     UpperCamelCase__ =			accelerator.prepare(
	    model, optimizer, train_dataloader, valid_dataloader, scheduler
	)
	UpperCamelCase__,     UpperCamelCase__ =			accelerator.prepare(model, optimizer)
	train(3, model, train_dataloader, optimizer, accelerator, scheduler)
	# Check that the intial optimizer is loaded on the GPU
	for group in optimizer.param_groups:
		UpperCamelCase__ =			group['params'][0].device
		break
	assert param_device.type == accelerator.device.type
	UpperCamelCase__ =			model.cpu()
	accelerator.wait_for_everyone()
	accelerator.save_state()
	accelerator.wait_for_everyone()

	# Check CPU state
	accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='cpu')
	for group in optimizer.param_groups:
		UpperCamelCase__ =			group['params'][0].device
		break
	assert (
	    param_device.type == torch.device('cpu').type
	), f"Loaded optimizer states did not match, expected to be loaded on the CPU but got {param_device}"

	# Check device state
	model.to(accelerator.device)
	accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='on_device')
	for group in optimizer.param_groups:
		UpperCamelCase__ =			group['params'][0].device
		break
	assert (
	    param_device.type == accelerator.device.type
	), f"Loaded optimizer states did not match, expected to be loaded on {accelerator.device} but got {param_device}"

	# Check error
	with pytest.raises(TypeError, match='Unsupported optimizer map location passed'):
		accelerator.load_state(os.path.join(savedir, 'checkpoints', 'checkpoint_0'), map_location='invalid')
	accelerator.wait_for_everyone()
	if accelerator.process_index == 0:
		shutil.rmtree(savedir)
	accelerator.wait_for_everyone()












'''simple docstring'''
import argparse
import os

import torch

from transformers import (
    XLNetConfig,
    XLNetForQuestionAnswering,
    XLNetForSequenceClassification,
    XLNetLMHeadModel,
    load_tf_weights_in_xlnet,
)
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


UpperCamelCase__ =			{
    'cola': 2,
    'mnli': 3,
    'mrpc': 2,
    'sst-2': 2,
    'sts-b': 1,
    'qqp': 2,
    'qnli': 2,
    'rte': 2,
    'wnli': 2,
}


logging.set_verbosity_info()




def 		__SCREAMING_SNAKE_CASE					(  _UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase=None							):


   """simple docstring"""




   lowercase_	:    Tuple           =     XLNetConfig.from_json_file(_UpperCamelCase							)

   lowercase_	:    Union[str, Any]           =     finetuning_task.lower() if finetuning_task is not None else ""
   if finetuning_task in GLUE_TASKS_NUM_LABELS:
      print(F"""Building PyTorch XLNetForSequenceClassification model from configuration: {config}"""							)
      lowercase_	:    Dict           =     finetuning_task
      lowercase_	:    Any           =     GLUE_TASKS_NUM_LABELS[finetuning_task]
      lowercase_	:    Any           =     XLNetForSequenceClassification(_UpperCamelCase							)
   elif "squad" in finetuning_task:
      lowercase_	:    Optional[int]           =     finetuning_task
      lowercase_	:    Optional[int]           =     XLNetForQuestionAnswering(_UpperCamelCase							)
   else:
      lowercase_	:    Union[str, Any]           =     XLNetLMHeadModel(_UpperCamelCase							)

   # Load weights from tf checkpoint
   load_tf_weights_in_xlnet(_UpperCamelCase						,		_UpperCamelCase						,		_UpperCamelCase							)

   # Save pytorch-model
   lowercase_	:    Optional[Any]           =     os.path.join(_UpperCamelCase						,		_UpperCamelCase							)
   lowercase_	:    Dict           =     os.path.join(_UpperCamelCase						,		_UpperCamelCase							)
   print(F"""Save PyTorch model to {os.path.abspath(_UpperCamelCase							)}"""							)
   torch.save(model.state_dict()						,		_UpperCamelCase							)
   print(F"""Save configuration file to {os.path.abspath(_UpperCamelCase							)}"""							)
   with open(_UpperCamelCase						,		"w"						,		encoding="utf-8"							) as f:
      f.write(config.to_json_string()							)


if __name__ == "__main__":
 UpperCamelCase__ =			argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '--tf_checkpoint_path', default=None, type=str, required=True, help='Path to the TensorFlow checkpoint path.'
 )
 parser.add_argument(
     '--xlnet_config_file',
     default=None,
     type=str,
     required=True,
     help=(
         'The config json file corresponding to the pre-trained XLNet model. \n'
         'This specifies the model architecture.'
     ),
 )
 parser.add_argument(
     '--pytorch_dump_folder_path',
     default=None,
     type=str,
     required=True,
     help='Path to the folder to store the PyTorch model or dataset/vocab.',
 )
 parser.add_argument(
     '--finetuning_task',
     default=None,
     type=str,
     help='Name of a task on which the XLNet TensorFlow model was fine-tuned',
 )
 UpperCamelCase__ =			parser.parse_args()
 print(args)

 convert_xlnet_checkpoint_to_pytorch(
     args.tf_checkpoint_path, args.xlnet_config_file, args.pytorch_dump_folder_path, args.finetuning_task
 )




