


import argparse
import collections
import os
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
snake_case							:       List[str]     			=      '''src/transformers'''
snake_case							:       Optional[int]     			=      '''docs/source/en'''
snake_case							:       Any     			=      '''.'''






def 							__lowercase   ( __lowerCAmelCase							:				List[Any]					,    __lowerCAmelCase							:				List[Any]					,    __lowerCAmelCase							:				Optional[int]   ):
							with open(__lowerCAmelCase					,    'r'					,    encoding='utf-8'					,    newline='\n'   ) as f:
														a__             =     f.readlines()
							# Find the start prompt.
							a__             =     0
							while not lines[start_index].startswith(__lowerCAmelCase   ):
														start_index += 1
							start_index += 1

							a__             =     start_index
							while not lines[end_index].startswith(__lowerCAmelCase   ):
														end_index += 1
							end_index -= 1

							while len(lines[start_index]   ) <= 1:
														start_index += 1
							while len(lines[end_index]   ) <= 1:
														end_index -= 1
							end_index += 1
							return "".join(lines[start_index:end_index]   ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
snake_case							:       Optional[int]     			=      '''Model|Encoder|Decoder|ForConditionalGeneration'''
# Regexes that match TF/Flax/PT model names.
snake_case							:       Dict     			=      re.compile(r'''TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')
snake_case							:       Union[str, Any]     			=      re.compile(r'''Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
snake_case							:       Any     			=      re.compile(r'''(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)''')


# This is to make sure the transformers module imported is the one in the repo.
snake_case							:       Optional[Any]     			=      direct_transformers_import(TRANSFORMERS_PATH)






def 							__lowercase   ( __lowerCAmelCase							:				Dict   ):
							a__             =     re.finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'					,    __lowerCAmelCase   )
							return [m.group(0   ) for m in matches]






def 							__lowercase   ( __lowerCAmelCase							:				Union[str, Any]					,    __lowerCAmelCase							:				Optional[int]   ):
							a__             =     2 if text == '✅' or text == '❌' else len(__lowerCAmelCase   )
							a__             =     (width - text_length) // 2
							a__             =     width - text_length - left_indent
							return " " * left_indent + text + " " * right_indent






def 							__lowercase   ( ):
							a__             =     transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
							a__             =     {
							    name: config_maping_names[code]
							    for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
							    if code in config_maping_names
							}
							a__             =     {name: config.replace('Config'					,    ''   ) for name, config in model_name_to_config.items()}

							# Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
							a__             =     collections.defaultdict(__lowerCAmelCase   )
							a__             =     collections.defaultdict(__lowerCAmelCase   )
							a__             =     collections.defaultdict(__lowerCAmelCase   )
							a__             =     collections.defaultdict(__lowerCAmelCase   )
							a__             =     collections.defaultdict(__lowerCAmelCase   )

							# Let's lookup through all transformers object (once).
							for attr_name in dir(__lowerCAmelCase   ):
														a__             =     None
														if attr_name.endswith('Tokenizer'   ):
																					a__             =     slow_tokenizers
																					a__             =     attr_name[:-9]
														elif attr_name.endswith('TokenizerFast'   ):
																					a__             =     fast_tokenizers
																					a__             =     attr_name[:-1_3]
														elif _re_tf_models.match(__lowerCAmelCase   ) is not None:
																					a__             =     tf_models
																					a__             =     _re_tf_models.match(__lowerCAmelCase   ).groups()[0]
														elif _re_flax_models.match(__lowerCAmelCase   ) is not None:
																					a__             =     flax_models
																					a__             =     _re_flax_models.match(__lowerCAmelCase   ).groups()[0]
														elif _re_pt_models.match(__lowerCAmelCase   ) is not None:
																					a__             =     pt_models
																					a__             =     _re_pt_models.match(__lowerCAmelCase   ).groups()[0]

														if lookup_dict is not None:
																					while len(__lowerCAmelCase   ) > 0:
																												if attr_name in model_name_to_prefix.values():
																																			a__             =     True
																																			break
																												# Try again after removing the last word in the name
																												a__             =     ''.join(camel_case_split(__lowerCAmelCase   )[:-1]   )

    # Let's build that table!
							a__             =     list(model_name_to_config.keys()   )
							model_names.sort(key=str.lower   )
							a__             =     ['Model', 'Tokenizer slow', 'Tokenizer fast', 'PyTorch support', 'TensorFlow support', 'Flax Support']
							# We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
							a__             =     [len(__lowerCAmelCase   ) + 2 for c in columns]
							a__             =     max([len(__lowerCAmelCase   ) for name in model_names]   ) + 2

							# Build the table per se
							a__             =     '|' + '|'.join([_center_text(__lowerCAmelCase					,    __lowerCAmelCase   ) for c, w in zip(__lowerCAmelCase					,    __lowerCAmelCase   )]   ) + '|\n'
							# Use ":-----:" format to center-aligned table cell texts
							table += "|" + "|".join([':' + '-' * (w - 2) + ':' for w in widths]   ) + "|\n"

							a__             =     {True: '✅', False: '❌'}
							for name in model_names:
														a__             =     model_name_to_prefix[name]
														a__             =     [
														    name,
														    check[slow_tokenizers[prefix]],
														    check[fast_tokenizers[prefix]],
														    check[pt_models[prefix]],
														    check[tf_models[prefix]],
														    check[flax_models[prefix]],
														]
														table += "|" + "|".join([_center_text(__lowerCAmelCase					,    __lowerCAmelCase   ) for l, w in zip(__lowerCAmelCase					,    __lowerCAmelCase   )]   ) + "|\n"
							return table






def 							__lowercase   ( __lowerCAmelCase							:				str=False   ):
							a__      ,					a__      ,					a__      ,					a__             =     _find_text_in_file(
							    filename=os.path.join(__lowerCAmelCase					,    'index.md'   )					,    start_prompt='<!--This table is updated automatically from the auto modules'					,    end_prompt='<!-- End table-->'					,    )
							a__             =     get_model_table_from_auto_modules()

							if current_table != new_table:
														if overwrite:
																					with open(os.path.join(__lowerCAmelCase					,    'index.md'   )					,    'w'					,    encoding='utf-8'					,    newline='\n'   ) as f:
																												f.writelines(lines[:start_index] + [new_table] + lines[end_index:]   )
														else:
																					raise ValueError(
																					    'The model table in the `index.md` has not been updated. Run `make fix-copies` to fix this.'   )


if __name__ == "__main__":
						snake_case							:       Optional[Any]     			=      argparse.ArgumentParser()
						parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
						snake_case							:       Union[str, Any]     			=      parser.parse_args()

						check_model_table(args.fix_and_overwrite)








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()




