
from sklearn.metrics import fa_score, matthews_corrcoef

import datasets

from .record_evaluation import evaluate as evaluate_record


SCREAMING_SNAKE_CASE							:Optional[int]							 =				'\\n@article{wang2019superglue,\n  title={SuperGLUE: A Stickier Benchmark for General-Purpose Language Understanding Systems},\n  author={Wang, Alex and Pruksachatkun, Yada and Nangia, Nikita and Singh, Amanpreet and Michael, Julian and Hill, Felix and Levy, Omer and Bowman, Samuel R},\n  journal={arXiv preprint arXiv:1905.00537},\n  year={2019}\n}\n'

SCREAMING_SNAKE_CASE							:Union[str, Any]							 =				'\\nSuperGLUE (https://super.gluebenchmark.com/) is a new benchmark styled after\nGLUE with a new set of more difficult language understanding tasks, improved\nresources, and a new public leaderboard.\n'

SCREAMING_SNAKE_CASE							:str							 =				'\nCompute SuperGLUE evaluation metric associated to each SuperGLUE dataset.\nArgs:\n    predictions: list of predictions to score. Depending on the SuperGlUE subset:\n        - for \'record\': list of question-answer dictionaries with the following keys:\n            - \'idx\': index of the question as specified by the dataset\n            - \'prediction_text\': the predicted answer text\n        - for \'multirc\': list of question-answer dictionaries with the following keys:\n            - \'idx\': index of the question-answer pair as specified by the dataset\n            - \'prediction\': the predicted answer label\n        - otherwise: list of predicted labels\n    references: list of reference labels. Depending on the SuperGLUE subset:\n        - for \'record\': list of question-answers dictionaries with the following keys:\n            - \'idx\': index of the question as specified by the dataset\n            - \'answers\': list of possible answers\n        - otherwise: list of reference labels\nReturns: depending on the SuperGLUE subset:\n    - for \'record\':\n        - \'exact_match\': Exact match between answer and gold answer\n        - \'f1\': F1 score\n    - for \'multirc\':\n        - \'exact_match\': Exact match between answer and gold answer\n        - \'f1_m\': Per-question macro-F1 score\n        - \'f1_a\': Average F1 score over all answers\n    - for \'axb\':\n        \'matthews_correlation\': Matthew Correlation\n    - for \'cb\':\n        - \'accuracy\': Accuracy\n        - \'f1\': F1 score\n    - for all others:\n        - \'accuracy\': Accuracy\nExamples:\n\n    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'copa\')  # any of ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]\n    >>> predictions = [0, 1]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'cb\')\n    >>> predictions = [0, 1]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'accuracy\': 1.0, \'f1\': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'record\')\n    >>> predictions = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'prediction_text\': \'answer\'}]\n    >>> references = [{\'idx\': {\'passage\': 0, \'query\': 0}, \'answers\': [\'answer\', \'another_answer\']}]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'exact_match\': 1.0, \'f1\': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'multirc\')\n    >>> predictions = [{\'idx\': {\'answer\': 0, \'paragraph\': 0, \'question\': 0}, \'prediction\': 0}, {\'idx\': {\'answer\': 1, \'paragraph\': 2, \'question\': 3}, \'prediction\': 1}]\n    >>> references = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'exact_match\': 1.0, \'f1_m\': 1.0, \'f1_a\': 1.0}\n\n    >>> super_glue_metric = datasets.load_metric(\'super_glue\', \'axb\')\n    >>> references = [0, 1]\n    >>> predictions = [0, 1]\n    >>> results = super_glue_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'matthews_correlation\': 1.0}\n'


def 						UpperCAmelCase			(	a_		,   a_						)		->	Dict:







				"""simple docstring"""



				return float((preds == labels).mean()						)


def 						UpperCAmelCase			(	a_		,   a_		,   a_="binary"						)		->	Union[str, Any]:







				"""simple docstring"""



				__A   			= simple_accuracy(a_		,   a_						)
				__A   			= float(fa_score(y_true=a_		,   y_pred=a_		,   average=a_						)						)
				return {
				    "accuracy": acc,
				    "f1": fa,
				}


def 						UpperCAmelCase			(	a_		,   a_						)		->	Union[str, Any]:







				"""simple docstring"""



				__A   			= {}
				for id_pred, label in zip(a_		,   a_						):
								__A   			= F'''{id_pred['idx']['paragraph']}-{id_pred['idx']['question']}'''
								__A   			= id_pred["prediction"]
								if question_id in question_map:
												question_map[question_id].append((pred, label)						)
								else:
												__A   			= [(pred, label)]
				__A			,							__A   			= [], []
				for question, preds_labels in question_map.items():
								__A			,							__A   			= zip(*a_						)
								__A   			= fa_score(y_true=a_		,   y_pred=a_		,   average="macro"						)
								fas.append(a_						)
								__A   			= int(sum(pred == label for pred, label in preds_labels						) == len(a_						)						)
								ems.append(a_						)
				__A   			= float(sum(a_						) / len(a_						)						)
				__A   			= sum(a_						) / len(a_						)
				__A   			= float(fa_score(y_true=a_		,   y_pred=[id_pred["prediction"] for id_pred in ids_preds]						)						)
				return {"exact_match": em, "f1_m": fa_m, "f1_a": fa_a}





@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION			,			_KWARGS_DESCRIPTION  )
class 							UpperCAmelCase      (  datasets.Metric  ):




	'''simple docstring'''





	def 				UpperCamelCase_  (       self					: List[str] ):
					if self.config_name not in [
					    "boolq",
					    "cb",
					    "copa",
					    "multirc",
					    "record",
					    "rte",
					    "wic",
					    "wsc",
					    "wsc.fixed",
					    "axb",
					    "axg",
					]:
									raise KeyError(
									    "You should supply a configuration name selected in "
									    "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]" )
					return datasets.MetricInfo(
					    description=_DESCRIPTION							,citation=_CITATION							,inputs_description=_KWARGS_DESCRIPTION							,features=datasets.Features(self._get_feature_types() )							,codebase_urls=[]							,reference_urls=[]							,format="numpy" if not self.config_name == "record" and not self.config_name == "multirc" else None							,)






	def 				UpperCamelCase_  (       self					: List[Any] ):
					if self.config_name == "record":
									return {
									    "predictions": {
									        "idx": {
									            "passage": datasets.Value("int64" ),
									            "query": datasets.Value("int64" ),
									        },
									        "prediction_text": datasets.Value("string" ),
									    },
									    "references": {
									        "idx": {
									            "passage": datasets.Value("int64" ),
									            "query": datasets.Value("int64" ),
									        },
									        "answers": datasets.Sequence(datasets.Value("string" ) ),
									    },
									}
					elif self.config_name == "multirc":
									return {
									    "predictions": {
									        "idx": {
									            "answer": datasets.Value("int64" ),
									            "paragraph": datasets.Value("int64" ),
									            "question": datasets.Value("int64" ),
									        },
									        "prediction": datasets.Value("int64" ),
									    },
									    "references": datasets.Value("int64" ),
									}
					else:
									return {
									    "predictions": datasets.Value("int64" ),
									    "references": datasets.Value("int64" ),
									}







	def 				UpperCamelCase_  (       self					: Union[str, Any]							,A					: Union[str, Any]							,A					: Dict ):
					if self.config_name == "axb":
									return {"matthews_correlation": matthews_corrcoef(A							,A )}
					elif self.config_name == "cb":
									return acc_and_fa(A							,A							,fa_avg="macro" )
					elif self.config_name == "record":
									__A   			= [
									    {
									        "qas": [
									            {"id": ref["idx"]["query"], "answers": [{"text": ans} for ans in ref["answers"]]}
									            for ref in references
									        ]
									    }
									]
									__A   			= {pred["idx"]["query"]: pred["prediction_text"] for pred in predictions}
									return evaluate_record(A							,A )[0]
					elif self.config_name == "multirc":
									return evaluate_multirc(A							,A )
					elif self.config_name in ["copa", "rte", "wic", "wsc", "wsc.fixed", "boolq", "axg"]:
									return {"accuracy": simple_accuracy(A							,A )}
					else:
									raise KeyError(
									    "You should supply a configuration name selected in "
									    "[\"boolq\", \"cb\", \"copa\", \"multirc\", \"record\", \"rte\", \"wic\", \"wsc\", \"wsc.fixed\", \"axb\", \"axg\",]" )




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


snake_case							:       List[Any]     			=      logging.getLogger(__name__)






def 							__lowercase   ( __lowerCAmelCase							:				str=2					,    __lowerCAmelCase							:				Optional[int]=3					,    __lowerCAmelCase							:				List[str]=1_6					,    __lowerCAmelCase							:				int = 1_0					,    __lowerCAmelCase							:				int = 2   ):

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

							a__             =     get_dataset(__lowerCAmelCase   )
							a__             =     get_dataset(__lowerCAmelCase   )
							a__             =     DataLoader(__lowerCAmelCase					,    shuffle=__lowerCAmelCase					,    batch_size=__lowerCAmelCase					,    num_workers=4   )
							a__             =     DataLoader(__lowerCAmelCase					,    shuffle=__lowerCAmelCase					,    batch_size=__lowerCAmelCase					,    num_workers=4   )
							return (train_dataloader, valid_dataloader)






def 							__lowercase   ( __lowerCAmelCase							:				Dict					,    __lowerCAmelCase							:				Union[str, Any]					,    __lowerCAmelCase							:				Tuple					,    __lowerCAmelCase							:				str					,    __lowerCAmelCase							:				int					,    __lowerCAmelCase							:				Optional[int]=None   ):
							a__             =     []
							for epoch in range(__lowerCAmelCase   ):
														# Train quickly
														model.train()
														for batch in dataloader:
																					a__      ,					a__             =     batch
																					a__             =     model(__lowerCAmelCase   )
																					a__             =     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        snake_case_   (nn.Module      ):
					def __init__(      self   :Any      )					->    Union[str, Any]:
												super().__init__()
												a__             =     nn.Parameter(torch.randn(1      )      )
												a__             =     nn.Parameter(torch.randn(1      )      )



					def  lowerCamelCase__(      self   :List[str]			,__snake_case   :Union[str, Any]      )					->    str:
												return x * self.a + self.b







class        snake_case_   (unittest.TestCase      ):
					def  lowerCamelCase__(      self   :Tuple      )					->    Dict:
												with tempfile.TemporaryDirectory() as tmpdir:
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__      ,					a__             =     dummy_dataloaders()
																			a__             =     ProjectConfiguration(total_limit=1			,project_dir=__snake_case			,automatic_checkpoint_naming=__snake_case      )
																			# Train baseline
																			a__             =     Accelerator(project_config=__snake_case      )
																			a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save initial
																			accelerator.save_state()

																			# Save second state
																			accelerator.save_state()
																			self.assertEqual(len(os.listdir(accelerator.project_dir      )      )			,1      )
					def  lowerCamelCase__(      self   :List[Any]      )					->    str:
												with tempfile.TemporaryDirectory() as tmpdir:
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__      ,					a__             =     dummy_dataloaders()
																			# Train baseline
																			a__             =     Accelerator()
																			a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save initial
																			a__             =     os.path.join(__snake_case			,'initial'      )
																			accelerator.save_state(__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			a__             =     train(3			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()

																			# Train partially
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__      ,					a__             =     dummy_dataloaders()
																			a__             =     Accelerator()
																			a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			accelerator.load_state(__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )

																			a__             =     train(2			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save everything
																			a__             =     os.path.join(__snake_case			,'checkpoint'      )
																			accelerator.save_state(__snake_case      )

																			# Load everything back in and make sure all states work
																			accelerator.load_state(__snake_case      )
																			test_rands += train(1			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
					def  lowerCamelCase__(      self   :str      )					->    List[str]:
												with tempfile.TemporaryDirectory() as tmpdir:
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__      ,					a__             =     dummy_dataloaders()
																			a__             =     ProjectConfiguration(automatic_checkpoint_naming=__snake_case      )

																			# Train baseline
																			a__             =     Accelerator(project_dir=__snake_case			,project_config=__snake_case      )
																			a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save initial
																			accelerator.save_state()
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			a__             =     train(3			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()

																			# Train partially
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__      ,					a__             =     dummy_dataloaders()
																			a__             =     ProjectConfiguration(iteration=1			,automatic_checkpoint_naming=__snake_case      )
																			a__             =     Accelerator(project_dir=__snake_case			,project_config=__snake_case      )
																			a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			accelerator.load_state(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_0'      )      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )

																			a__             =     train(2			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save everything
																			accelerator.save_state()

																			# Load everything back in and make sure all states work
																			accelerator.load_state(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_1'      )      )
																			test_rands += train(1			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			((a__)      ,					(a__))             =     model.a.item(), model.b.item()
																			a__             =     optimizer.state_dict()
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
																			self.assertEqual(__snake_case			,__snake_case      )
					def  lowerCamelCase__(      self   :Union[str, Any]      )					->    List[str]:
												a__             =     torch.tensor([1, 2, 3]      )
												a__             =     torch.tensor([2, 3, 4]      )
												a__             =     DummyModel()
												a__             =     torch.optim.Adam(net.parameters()      )
												a__             =     Accelerator()
												with self.assertRaises(__snake_case      ) as ve:
																			accelerator.register_for_checkpointing(__snake_case			,__snake_case			,__snake_case			,__snake_case      )
												a__             =     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[Any]      )					->    Dict:
												with tempfile.TemporaryDirectory() as tmpdir:
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     torch.optim.Adam(params=model.parameters()			,lr=1E-3      )
																			a__             =     torch.optim.lr_scheduler.StepLR(__snake_case			,step_size=1			,gamma=0.99      )
																			a__      ,					a__             =     dummy_dataloaders()
																			a__             =     ProjectConfiguration(automatic_checkpoint_naming=__snake_case      )
																			# Train baseline
																			a__             =     Accelerator(project_dir=__snake_case			,project_config=__snake_case      )
																			a__      ,					a__      ,					a__      ,					a__      ,					a__             =     accelerator.prepare(
																			    __snake_case			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			# Save initial
																			accelerator.save_state()
																			a__             =     scheduler.state_dict()
																			train(3			,__snake_case			,__snake_case			,__snake_case			,__snake_case			,__snake_case      )
																			self.assertNotEqual(__snake_case			,scheduler.state_dict()      )

																			# Load everything back in and make sure all states work
																			accelerator.load_state(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_0'      )      )
																			self.assertEqual(__snake_case			,scheduler.state_dict()      )
					def  lowerCamelCase__(      self   :Optional[int]      )					->    List[str]:
												with tempfile.TemporaryDirectory() as tmpdir:
																			set_seed(42      )
																			a__             =     DummyModel()
																			a__             =     ProjectConfiguration(automatic_checkpoint_naming=__snake_case			,total_limit=2      )
																			# Train baseline
																			a__             =     Accelerator(project_dir=__snake_case			,project_config=__snake_case      )
																			a__             =     accelerator.prepare(__snake_case      )
																			# Save 3 states:
																			for _ in range(11      ):
																										accelerator.save_state()
																			self.assertTrue(not os.path.exists(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_0'      )      )      )
																			self.assertTrue(os.path.exists(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_9'      )      )      )
																			self.assertTrue(os.path.exists(os.path.join(__snake_case			,'checkpoints'			,'checkpoint_10'      )      )      )



					@require_cuda
					def  lowerCamelCase__(      self   :Dict      )					->    str:
												a__             =     ['torchrun', F'--nproc_per_node={torch.cuda.device_count()}', inspect.getfile(self.__class__      )]
												execute_subprocess_async(__snake_case			,env=os.environ.copy()      )


if __name__ == "__main__":
						snake_case							:       Tuple     			=      '''/tmp/accelerate/state_checkpointing'''
						snake_case							:       str     			=      DummyModel()
						snake_case							:       List[Any]     			=      torch.optim.Adam(params=model.parameters(), lr=1e-3)
						snake_case							:       Union[str, Any]     			=      torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.99)
						snake_case			,     snake_case							:       str     			=      dummy_dataloaders()
						snake_case							:       Union[str, Any]     			=      ProjectConfiguration(automatic_checkpoint_naming=True)
						# Train baseline
						snake_case							:       Dict     			=      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)
						snake_case			,     snake_case			,     snake_case			,     snake_case			,     snake_case							:       List[str]     			=      accelerator.prepare(
						    model, optimizer, train_dataloader, valid_dataloader, scheduler
						)
						snake_case			,     snake_case							:       Any     			=      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:
												snake_case							:       Any     			=      group['''params'''][0].device
												break
						assert param_device.type == accelerator.device.type
						snake_case							:       Union[str, Any]     			=      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:
												snake_case							:       int     			=      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:
												snake_case							:       Optional[int]     			=      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()




