


"""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


__snake_case		=       logging.getLogger(__name__)
def       A_   (   _lowerCAmelCase     :					str=2,     _lowerCAmelCase     :					Optional[int]=3,     _lowerCAmelCase     :					Optional[Any]=16,     _lowerCAmelCase     :					int = 10,     _lowerCAmelCase     :					int = 2       ):






     """simple docstring"""

     def get_dataset(_lowerCAmelCase     :					Optional[Any]       ):
          _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       A_   (   _lowerCAmelCase     :					Any,     _lowerCAmelCase     :					Tuple,     _lowerCAmelCase     :					int,     _lowerCAmelCase     :					Union[str, Any],     _lowerCAmelCase     :					Union[str, Any],     _lowerCAmelCase     :					Union[str, Any]=None       ):






     """simple docstring"""

     _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     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self							)       ->							List[Any]:
          super().__init__()
          _a    	=				nn.Parameter(torch.randn(1							)							)
          _a    	=				nn.Parameter(torch.randn(1							)							)





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







class     __lowerCamelCase  (					unittest.TestCase				):

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self							)       ->							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(total_limit=1    ,				project_dir=__UpperCAmelCase    ,				automatic_checkpoint_naming=__UpperCAmelCase							)
               # Train baseline
               _a    	=				Accelerator(project_config=__UpperCAmelCase							)
               _a	,      _a	,      _a	,      _a    	=				accelerator.prepare(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               # Save initial
               accelerator.save_state()

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


     def 							_UpperCAmelCase      (      self							)       ->							Tuple:
          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(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               # Save initial
               _a    	=				os.path.join(__UpperCAmelCase    ,				'''initial'''							)
               accelerator.save_state(__UpperCAmelCase							)
               ((_a)	,      (_a))    	=				model.a.item(), model.b.item()
               _a    	=				optimizer.state_dict()
               _a    	=				train(3    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               ((_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(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               accelerator.load_state(__UpperCAmelCase							)
               ((_a)	,      (_a))    	=				model.a.item(), model.b.item()
               _a    	=				optimizer.state_dict()
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)

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

               # Load everything back in and make sure all states work
               accelerator.load_state(__UpperCAmelCase							)
               test_rands += train(1    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               ((_a)	,      (_a))    	=				model.a.item(), model.b.item()
               _a    	=				optimizer.state_dict()
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)


     def 							_UpperCAmelCase      (      self							)       ->							Optional[int]:
          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=__UpperCAmelCase							)

               # Train baseline
               _a    	=				Accelerator(project_dir=__UpperCAmelCase    ,				project_config=__UpperCAmelCase							)
               _a	,      _a	,      _a	,      _a    	=				accelerator.prepare(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               # Save initial
               accelerator.save_state()
               ((_a)	,      (_a))    	=				model.a.item(), model.b.item()
               _a    	=				optimizer.state_dict()
               _a    	=				train(3    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               ((_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=__UpperCAmelCase							)
               _a    	=				Accelerator(project_dir=__UpperCAmelCase    ,				project_config=__UpperCAmelCase							)
               _a	,      _a	,      _a	,      _a    	=				accelerator.prepare(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               accelerator.load_state(os.path.join(__UpperCAmelCase    ,				'''checkpoints'''    ,				'''checkpoint_0'''							)							)
               ((_a)	,      (_a))    	=				model.a.item(), model.b.item()
               _a    	=				optimizer.state_dict()
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertEqual(__UpperCAmelCase    ,				__UpperCAmelCase							)

               _a    	=				train(2    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               # Save everything
               accelerator.save_state()

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


     def 							_UpperCAmelCase      (      self							)       ->							Union[str, Any]:
          _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(__UpperCAmelCase							) as ve:
               accelerator.register_for_checkpointing(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
          _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 							_UpperCAmelCase      (      self							)       ->							Any:
          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(__UpperCAmelCase    ,				step_size=1    ,				gamma=0.99							)
               _a	,      _a    	=				dummy_dataloaders()
               _a    	=				ProjectConfiguration(automatic_checkpoint_naming=__UpperCAmelCase							)
               # Train baseline
               _a    	=				Accelerator(project_dir=__UpperCAmelCase    ,				project_config=__UpperCAmelCase							)
               _a	,      _a	,      _a	,      _a	,      _a    	=				accelerator.prepare(
                   __UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               # Save initial
               accelerator.save_state()
               _a    	=				scheduler.state_dict()
               train(3    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
               self.assertNotEqual(__UpperCAmelCase    ,				scheduler.state_dict()							)

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


     def 							_UpperCAmelCase      (      self							)       ->							List[str]:
          with tempfile.TemporaryDirectory() as tmpdir:
               set_seed(42							)
               _a    	=				DummyModel()
               _a    	=				ProjectConfiguration(automatic_checkpoint_naming=__UpperCAmelCase    ,				total_limit=2							)
               # Train baseline
               _a    	=				Accelerator(project_dir=__UpperCAmelCase    ,				project_config=__UpperCAmelCase							)
               _a    	=				accelerator.prepare(__UpperCAmelCase							)
               # Save 3 states:
               for _ in range(11							):
                    accelerator.save_state()
               self.assertTrue(not os.path.exists(os.path.join(__UpperCAmelCase    ,				'''checkpoints'''    ,				'''checkpoint_0'''							)							)							)
               self.assertTrue(os.path.exists(os.path.join(__UpperCAmelCase    ,				'''checkpoints'''    ,				'''checkpoint_9'''							)							)							)
               self.assertTrue(os.path.exists(os.path.join(__UpperCAmelCase    ,				'''checkpoints'''    ,				'''checkpoint_10'''							)							)							)





     @require_cuda
     def 							_UpperCAmelCase      (      self							)       ->							List[Any]:
          _a    	=				['''torchrun''', F'--nproc_per_node={torch.cuda.device_count()}', inspect.getfile(self.__class__							)]
          execute_subprocess_async(__UpperCAmelCase    ,				env=os.environ.copy()							)


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

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


__snake_case		=       logging.get_logger(__name__)


class     __lowerCamelCase  (					a__				):

     '''simple docstring'''
     A_       :	Any    =		['input_ids', 'attention_mask']


     def __init__(      self    ,				__UpperCAmelCase="</s>"    ,				__UpperCAmelCase="<unk>"    ,				__UpperCAmelCase="<pad>"    ,				__UpperCAmelCase=125    ,				__UpperCAmelCase=None    ,				**__UpperCAmelCase    ,				)       ->							None:
          # Add extra_ids to the special token list
          if extra_ids > 0 and additional_special_tokens is None:
               _a    	=				[F'<extra_id_{i}>' for i in range(__UpperCAmelCase							)]
          elif extra_ids > 0 and additional_special_tokens is not None:
               # Check that we have the right number of extra_id special tokens
               _a    	=				len(set(filter(lambda __UpperCAmelCase							: bool('''extra_id''' in str(__UpperCAmelCase							)							)    ,				__UpperCAmelCase							)							)							)
               if extra_tokens != extra_ids:
                    raise ValueError(
                        F'Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are'
                        ''' provided to ByT5Tokenizer. In this case the additional_special_tokens must include the'''
                        ''' extra_ids tokens'''							)

          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else pad_token
          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else eos_token
          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else unk_token

          super().__init__(
              eos_token=__UpperCAmelCase    ,				unk_token=__UpperCAmelCase    ,				pad_token=__UpperCAmelCase    ,				extra_ids=__UpperCAmelCase    ,				additional_special_tokens=__UpperCAmelCase    ,				**__UpperCAmelCase    ,				)

          _a    	=				extra_ids

          _a    	=				2**8  # utf is 8 bits

          # define special tokens dict
          _a    	=				{
              self.pad_token: 0,
              self.eos_token: 1,
              self.unk_token: 2,
          }
          _a    	=				len(self.special_tokens_encoder							)
          _a    	=				len(__UpperCAmelCase							)
          for i, token in enumerate(__UpperCAmelCase							):
               _a    	=				self.vocab_size + i - n
          _a    	=				{v: k for k, v in self.special_tokens_encoder.items()}


     @property
     def 							_UpperCAmelCase      (      self							)       ->							Dict:
          return self._utf_vocab_size + self._num_special_tokens + self._extra_ids


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = False							)       ->							List[int]:
          if already_has_special_tokens:
               return super().get_special_tokens_mask(
                   token_ids_a=__UpperCAmelCase    ,				token_ids_a=__UpperCAmelCase    ,				already_has_special_tokens=__UpperCAmelCase							)

          # normal case: some special tokens
          if token_ids_a is None:
               return ([0] * len(__UpperCAmelCase							)) + [1]
          return ([0] * len(__UpperCAmelCase							)) + [1] + ([0] * len(__UpperCAmelCase							)) + [1]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[int]:
          if len(__UpperCAmelCase							) > 0 and token_ids[-1] == self.eos_token_id:
               warnings.warn(
                   F'This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated'
                   ''' eos tokens being added.'''							)
               return token_ids
          else:
               return token_ids + [self.eos_token_id]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.eos_token_id]

          if token_ids_a is None:
               return len(token_ids_a + eos							) * [0]
          return len(token_ids_a + eos + token_ids_a + eos							) * [0]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				self._add_eos_if_not_present(__UpperCAmelCase							)
          if token_ids_a is None:
               return token_ids_a
          else:
               _a    	=				self._add_eos_if_not_present(__UpperCAmelCase							)
               return token_ids_a + token_ids_a


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          _a    	=				[chr(__UpperCAmelCase							) for i in text.encode('''utf-8'''							)]
          return tokens


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Union[str, Any]:
          if token in self.special_tokens_encoder:
               _a    	=				self.special_tokens_encoder[token]
          elif token in self.added_tokens_encoder:
               _a    	=				self.added_tokens_encoder[token]
          elif len(__UpperCAmelCase							) != 1:
               _a    	=				self.unk_token_id
          else:
               _a    	=				ord(__UpperCAmelCase							) + self._num_special_tokens
          return token_id


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[Any]:
          if index in self.special_tokens_decoder:
               _a    	=				self.special_tokens_decoder[index]
          else:
               _a    	=				chr(index - self._num_special_tokens							)
          return token


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Any:
          _a    	=				B''''''
          for token in tokens:
               if token in self.special_tokens_decoder:
                    _a    	=				self.special_tokens_decoder[token].encode('''utf-8'''							)
               elif token in self.added_tokens_decoder:
                    _a    	=				self.special_tokens_decoder[token].encode('''utf-8'''							)
               elif token in self.special_tokens_encoder:
                    _a    	=				token.encode('''utf-8'''							)
               elif token in self.added_tokens_encoder:
                    _a    	=				token.encode('''utf-8'''							)
               else:
                    _a    	=				bytes([ord(__UpperCAmelCase							)]							)
               bstring += tok_string
          _a    	=				bstring.decode('''utf-8'''    ,				errors='''ignore'''							)
          return string





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							Tuple[str]:
          return ()