


'''simple docstring'''






import json
import os
import tempfile
from unittest.mock import patch

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

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment






def 	_A				()   ->  Dict:




     '''simple docstring'''



     _a			      =							torch.nn.Linear(2				,						4		)
     _a			      =							torch.optim.AdamW(model.parameters()				,						lr=1.0		)
     _a			      =							torch.optim.lr_scheduler.OneCycleLR(lowerCAmelCase__				,						max_lr=0.0_1				,						steps_per_epoch=2				,						epochs=1		)
     _a			      =							DataLoader(TensorDataset(torch.tensor([1, 2, 3]		)		)		)
     _a			      =							DataLoader(TensorDataset(torch.tensor([4, 5, 6]		)		)		)

     return model, optimizer, scheduler, train_dl, valid_dl






def 	_A				(lowerCAmelCase__  :str		)   ->  Optional[int]:




     '''simple docstring'''



     return (model.weight.abs().sum() + model.bias.abs().sum()).item()






def 	_A				(lowerCAmelCase__  :Optional[Any]		)   ->  Optional[int]:




     '''simple docstring'''



     _a			      =							torch.nn.Linear(*tuple(model.weight.T.shape		)		).state_dict()
     model.load_state_dict(lowerCAmelCase__		)






class  a     (     _SCREAMING_SNAKE_CASE       ):
 @require_cuda
 def 							__UpperCAmelCase		( self						)					->							Union[str, Any]:
      _a			      =							Accelerator()
      assert PartialState._shared_state["_cpu"] is False
      assert PartialState._shared_state["device"].type == "cuda"
      with self.assertRaises(__magic_name__						):
           _a			      =							Accelerator(cpu=__magic_name__						)
 def 							__UpperCAmelCase		( self						)					->							Optional[int]:
      _a			      =							Accelerator()
      _a			      =							GradientState()
      assert state.num_steps == 1
      _a			      =							4
      assert state.num_steps == 4

      assert state.sync_gradients is True
      _a			      =							False
      assert state.sync_gradients is False
      GradientState._reset_state()
 def 							__UpperCAmelCase		( self						)					->							List[str]:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()

      (
          (
          _a   
      )      ,   (
          _a   
      )      ,   (
          _a   
      )      ,   (
          _a   
      )      ,   (
          _a   
      )      ,      
      )			      =							accelerator.prepare(__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)

      self.assertTrue(prepared_model in accelerator._models						)
      self.assertTrue(prepared_optimizer in accelerator._optimizers						)
      self.assertTrue(prepared_scheduler in accelerator._schedulers						)
      self.assertTrue(prepared_train_dl in accelerator._dataloaders						)
      self.assertTrue(prepared_valid_dl in accelerator._dataloaders						)
 def 							__UpperCAmelCase		( self						)					->							Tuple:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()
      accelerator.prepare(__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)
      accelerator.free_memory()

      self.assertTrue(len(accelerator._models						) == 0						)
      self.assertTrue(len(accelerator._optimizers						) == 0						)
      self.assertTrue(len(accelerator._schedulers						) == 0						)
      self.assertTrue(len(accelerator._dataloaders						) == 0						)
 def 							__UpperCAmelCase		( self						)					->							Union[str, Any]:
      PartialState._reset_state()

      # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
      def noop(*__magic_name__      ,				**__magic_name__						):
           pass

      with patch('torch.cuda.set_device'      ,				__magic_name__						), patch_environment(ACCELERATE_TORCH_DEVICE='cuda:64'						):
           _a			      =							Accelerator()
           self.assertEqual(str(accelerator.state.device						)      ,				'cuda:64'						)
 def 							__UpperCAmelCase		( self						)					->							List[str]:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()
      accelerator.prepare(__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)

      _a			      =							get_signature(__magic_name__						)

      with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(__magic_name__						)

           # make sure random weights don't match
           load_random_weights(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) > 1e-3						)

           # make sure loaded weights match
           accelerator.load_state(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) < 1e-3						)
 def 							__UpperCAmelCase		( self						)					->							Tuple:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()
      accelerator.prepare(__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)

      _a			      =							get_signature(__magic_name__						)

      # saving hook
      def save_config(__magic_name__      ,				__magic_name__      ,				__magic_name__						):
           _a			      =							{'class_name': models[0].__class__.__name__}

           with open(os.path.join(__magic_name__      ,				'data.json'						)      ,				'w'						) as f:
                json.dump(__magic_name__      ,				__magic_name__						)

        # loading hook
      def load_config(__magic_name__      ,				__magic_name__						):
           with open(os.path.join(__magic_name__      ,				'data.json'						)      ,				'r'						) as f:
                _a			      =							json.load(__magic_name__						)

           _a			      =							config['class_name']

      _a			      =							accelerator.register_save_state_pre_hook(__magic_name__						)
      _a			      =							accelerator.register_load_state_pre_hook(__magic_name__						)

      with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(__magic_name__						)

           # make sure random weights don't match with hooks
           load_random_weights(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) > 1e-3						)

           # random class name to verify correct one is loaded
           _a			      =							'random'

           # make sure loaded weights match with hooks
           accelerator.load_state(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) < 1e-3						)

           # mode.class_name is loaded from config
           self.assertTrue(model.class_name == model.__class__.__name__						)

      # remove hooks
      save_hook.remove()
      load_hook.remove()

      with tempfile.TemporaryDirectory() as tmpdirname:
           accelerator.save_state(__magic_name__						)

           # make sure random weights don't match with hooks removed
           load_random_weights(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) > 1e-3						)

           # random class name to verify correct one is loaded
           _a			      =							'random'

           # make sure loaded weights match with hooks removed
           accelerator.load_state(__magic_name__						)
           self.assertTrue(abs(model_signature - get_signature(__magic_name__						)						) < 1e-3						)

           # mode.class_name is NOT loaded from config
           self.assertTrue(model.class_name != model.__class__.__name__						)
 def 							__UpperCAmelCase		( self						)					->							List[str]:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()
      _a			      =							None

      # This should work
      _a      ,   _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							accelerator.prepare(
          __magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)
      self.assertTrue(dummy_obj is None						)
 def 							__UpperCAmelCase		( self						)					->							str:
      _a			      =							Accelerator()
      _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							create_components()
      _a			      =							[1, 2, 3]

      # This should work
      _a      ,   _a      ,   _a      ,   _a      ,   _a      ,   _a			      =							accelerator.prepare(
          __magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__      ,				__magic_name__						)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Dummy object should have `_is_accelerate_prepared` set to `True`'      ,				)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Model is missing `_is_accelerator_prepared` or is set to `False`'      ,				)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Optimizer is missing `_is_accelerator_prepared` or is set to `False`'      ,				)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Scheduler is missing `_is_accelerator_prepared` or is set to `False`'      ,				)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'      ,				)
      self.assertEqual(
          getattr(__magic_name__      ,				'_is_accelerate_prepared'      ,				__magic_name__						)      ,				__magic_name__      ,				'Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'      ,				)
 @slow
 @require_bnb
 def 							__UpperCAmelCase		( self						)					->							Optional[int]:
      from transformers import AutoModelForCausalLM

      _a			      =							AutoModelForCausalLM.from_pretrained(
          'EleutherAI/gpt-neo-125m'      ,				load_in_abit=__magic_name__      ,				device_map={'': 0}      ,				)
      _a			      =							Accelerator()

      # This should work
      _a			      =							accelerator.prepare(__magic_name__						)
 @slow
 @require_bnb
 def 							__UpperCAmelCase		( self						)					->							Optional[int]:
      from transformers import AutoModelForCausalLM

      _a			      =							Accelerator()

      with init_empty_weights():
           _a			      =							AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m'      ,				)
           model.tie_weights()
           _a			      =							infer_auto_device_map(__magic_name__						)
           _a			      =							'cpu'

      _a			      =							AutoModelForCausalLM.from_pretrained(
          'EleutherAI/gpt-neo-125m'      ,				device_map=__magic_name__      ,				load_in_abit=__magic_name__      ,				llm_inta_enable_fpaa_cpu_offload=__magic_name__						)

      # This should not work and get value error
      with self.assertRaises(__magic_name__						):
           _a			      =							accelerator.prepare(__magic_name__						)
 @slow
 @require_bnb
 @require_multi_gpu
 def 							__UpperCAmelCase		( self						)					->							List[Any]:
      from transformers import AutoModelForCausalLM

      _a			      =							{'distributed_type': DistributedType.MULTI_GPU}

      with init_empty_weights():
           _a			      =							AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m'      ,				)
           model.tie_weights()
           _a			      =							infer_auto_device_map(__magic_name__						)
           _a			      =							1

      _a			      =							AutoModelForCausalLM.from_pretrained(
          'EleutherAI/gpt-neo-125m'      ,				load_in_abit=__magic_name__      ,				device_map=__magic_name__      ,				)
      _a			      =							Accelerator()

      # This should not work and get value error
      with self.assertRaises(__magic_name__						):
           _a			      =							accelerator.prepare(__magic_name__						)

      PartialState._reset_state()
 @slow
 @require_bnb
 @require_multi_gpu
 def 							__UpperCAmelCase		( self						)					->							Optional[int]:
      from transformers import AutoModelForCausalLM

      with init_empty_weights():
           _a			      =							AutoModelForCausalLM.from_pretrained(
               'EleutherAI/gpt-neo-125m'      ,				)
           _a			      =							infer_auto_device_map(__magic_name__						)
           _a			      =							1

      _a			      =							AutoModelForCausalLM.from_pretrained(
          'EleutherAI/gpt-neo-125m'      ,				load_in_abit=__magic_name__      ,				device_map=__magic_name__      ,				)
      _a			      =							Accelerator()

      # This should work
      _a			      =							accelerator.prepare(__magic_name__						)




 @require_cuda
 def 							__UpperCAmelCase		( self						)					->							List[Any]:
      _a			      =							torch.nn.Linear(10      ,				10						)
      _a			      =							torch.optim.SGD(model.parameters()      ,				lr=0.0_1						)
      _a			      =							Accelerator(cpu=__magic_name__						)
      _a			      =							accelerator.prepare(__magic_name__						)





'''simple docstring'''






from typing import List, Optional

import numpy as np

from ...processing_utils import ProcessorMixin
from ...utils import to_numpy




class  a     (     _SCREAMING_SNAKE_CASE       ):
 _lowerCAmelCase 					=	"""EncodecFeatureExtractor"""
 _lowerCAmelCase 					=	("""T5Tokenizer""", """T5TokenizerFast""")
 def __init__( self      ,				__magic_name__      ,				__magic_name__						)					->							Optional[Any]:
      super().__init__(__magic_name__      ,				__magic_name__						)
      _a			      =							self.feature_extractor
      _a			      =							False
 def 							__UpperCAmelCase		( self      ,				__magic_name__=None      ,				__magic_name__=None      ,				__magic_name__=True						)					->							Tuple:
      return self.tokenizer.get_decoder_prompt_ids(task=__magic_name__      ,				language=__magic_name__      ,				no_timestamps=__magic_name__						)
 def __call__( self      ,				*__magic_name__      ,				**__magic_name__						)					->							List[str]:
      # For backward compatibility
      if self._in_target_context_manager:
           return self.current_processor(*__magic_name__      ,				**__magic_name__						)

      _a			      =							kwargs.pop('audio'      ,				__magic_name__						)
      _a			      =							kwargs.pop('sampling_rate'      ,				__magic_name__						)
      _a			      =							kwargs.pop('text'      ,				__magic_name__						)
      if len(__magic_name__						) > 0:
           _a			      =							args[0]
           _a			      =							args[1:]

      if audio is None and text is None:
           raise ValueError('You need to specify either an `audio` or `text` input to process.'						)

      if text is not None:
           _a			      =							self.tokenizer(__magic_name__      ,				**__magic_name__						)

      if audio is not None:
           _a			      =							self.feature_extractor(__magic_name__      ,				*__magic_name__      ,				sampling_rate=__magic_name__      ,				**__magic_name__						)

      if audio is None:
           return inputs

      elif text is None:
           return audio_inputs

      else:
           _a			      =							audio_inputs['input_values']
           if "padding_mask" in audio_inputs:
                _a			      =							audio_inputs['padding_mask']
           return inputs
 def 							__UpperCAmelCase		( self      ,				*__magic_name__      ,				**__magic_name__						)					->							Tuple:
      _a			      =							kwargs.pop('audio'      ,				__magic_name__						)
      _a			      =							kwargs.pop('padding_mask'      ,				__magic_name__						)

      if len(__magic_name__						) > 0:
           _a			      =							args[0]
           _a			      =							args[1:]

      if audio_values is not None:
           return self._decode_audio(__magic_name__      ,				padding_mask=__magic_name__						)
      else:
           return self.tokenizer.batch_decode(*__magic_name__      ,				**__magic_name__						)
 def 							__UpperCAmelCase		( self      ,				*__magic_name__      ,				**__magic_name__						)					->							List[Any]:
      return self.tokenizer.decode(*__magic_name__      ,				**__magic_name__						)




 def 							__UpperCAmelCase		( self      ,				__magic_name__      ,				__magic_name__ = None						)					->							List[np.ndarray]:
      _a			      =							to_numpy(__magic_name__						)
      _a      ,   _a      ,   _a			      =							audio_values.shape

      if padding_mask is None:
           return list(__magic_name__						)

      _a			      =							to_numpy(__magic_name__						)

      # match the sequence length of the padding mask to the generated audio arrays by padding with the **non-padding**
      # token (so that the generated audio values are **not** treated as padded tokens)
      _a			      =							seq_len - padding_mask.shape[-1]
      _a			      =							1 - self.feature_extractor.padding_value
      _a			      =							np.pad(__magic_name__      ,				((0, 0), (0, difference))      ,				'constant'      ,				constant_values=__magic_name__						)

      _a			      =							audio_values.tolist()
      for i in range(__magic_name__						):
           _a			      =							np.asarray(audio_values[i]						)[
               padding_mask[i][None, :] != self.feature_extractor.padding_value
           ]
           _a			      =							sliced_audio.reshape(__magic_name__      ,				-1						)

      return audio_values

