




import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
       from transformers import SpeechaTextFeatureExtractor

UpperCAmelCase_									=	random.Random()




def 					lowerCamelCase__       (						A__       :					Any  ,			A__       :					Union[str, Any]=1.0  ,			A__       :					int=None  ,			A__       :					Union[str, Any]=None	):





       '''simple docstring'''



       if rng is None:
              __lowerCamelCase   =						global_rng

       __lowerCamelCase   =						[]
       for batch_idx in range(shape[0]	):
              values.append([]	)
              for _ in range(shape[1]	):
                     values[-1].append(rng.random() * scale	)

       return values





@require_torch
@require_torchaudio
class     lowerCamelCase__(					unittest.TestCase):
      def __init__(				self: Tuple    ,       UpperCamelCase_: int    ,       UpperCamelCase_: Union[str, Any]=7    ,       UpperCamelCase_: List[str]=4_00    ,       UpperCamelCase_: Union[str, Any]=20_00    ,       UpperCamelCase_: Union[str, Any]=24    ,       UpperCamelCase_: List[Any]=24    ,       UpperCamelCase_: Union[str, Any]=0.0    ,       UpperCamelCase_: List[str]=1_60_00    ,       UpperCamelCase_: Union[str, Any]=True    ,       UpperCamelCase_: Union[str, Any]=True    ,       ):
             __lowerCamelCase   =						parent
             __lowerCamelCase   =						batch_size
             __lowerCamelCase   =						min_seq_length
             __lowerCamelCase   =						max_seq_length
             __lowerCamelCase   =						(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
             __lowerCamelCase   =						feature_size
             __lowerCamelCase   =						num_mel_bins
             __lowerCamelCase   =						padding_value
             __lowerCamelCase   =						sampling_rate
             __lowerCamelCase   =						return_attention_mask
             __lowerCamelCase   =						do_normalize
      def 			lowerCAmelCase__						(				self: Optional[int]   ):
             return {
                 "feature_size": self.feature_size,
                 "num_mel_bins": self.num_mel_bins,
                 "padding_value": self.padding_value,
                 "sampling_rate": self.sampling_rate,
                 "return_attention_mask": self.return_attention_mask,
                 "do_normalize": self.do_normalize,
             }




      def 			lowerCAmelCase__						(				self: str    ,       UpperCamelCase_: List[Any]=False    ,       UpperCamelCase_: Tuple=False   ):
             def _flatten(UpperCamelCase_: int   ):
                    return list(itertools.chain(*UpperCamelCase_   )   )

             if equal_length:
                    __lowerCamelCase   =						[floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
             else:
                    # make sure that inputs increase in size
                    __lowerCamelCase   =						[
                        floats_list((x, self.feature_size)   )
                        for x in range(self.min_seq_length    ,       self.max_seq_length    ,       self.seq_length_diff   )
                    ]
             if numpify:
                    __lowerCamelCase   =						[np.asarray(UpperCamelCase_   ) for x in speech_inputs]
             return speech_inputs





@require_torch
@require_torchaudio
class     lowerCamelCase__(					__lowerCamelCase ,     unittest.TestCase):
      UpperCAmelCase__      :	str    		=   SpeechaTextFeatureExtractor if is_speech_available() else None
      def 			lowerCAmelCase__						(				self: Tuple   ):
             __lowerCamelCase   =						SpeechaTextFeatureExtractionTester(self   )
      def 			lowerCAmelCase__						(				self: Tuple    ,       UpperCamelCase_: Any   ):
             self.assertTrue(np.all(np.mean(UpperCamelCase_    ,       axis=0   ) < 1E-3   )   )
             self.assertTrue(np.all(np.abs(np.var(UpperCamelCase_    ,       axis=0   ) - 1   ) < 1E-3   )   )
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             # Tests that all call wrap to encode_plus and batch_encode_plus
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             # create three inputs of length 800, 1000, and 1200
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]
             __lowerCamelCase   =						[np.asarray(UpperCamelCase_   ) for speech_input in speech_inputs]

             # Test feature size
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       padding=UpperCamelCase_    ,       return_tensors="""np"""   ).input_features
             self.assertTrue(input_features.ndim == 3   )
             self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size   )

             # Test not batched input
             __lowerCamelCase   =						feature_extractor(speech_inputs[0]    ,       return_tensors="""np"""   ).input_features
             __lowerCamelCase   =						feature_extractor(np_speech_inputs[0]    ,       return_tensors="""np"""   ).input_features
             self.assertTrue(np.allclose(UpperCamelCase_    ,       UpperCamelCase_    ,       atol=1E-3   )   )

             # Test batched
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       return_tensors="""np"""   ).input_features
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       return_tensors="""np"""   ).input_features
             for enc_seq_a, enc_seq_a in zip(UpperCamelCase_    ,       UpperCamelCase_   ):
                    self.assertTrue(np.allclose(UpperCamelCase_    ,       UpperCamelCase_    ,       atol=1E-3   )   )

             # Test 2-D numpy arrays are batched.
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in (8_00, 8_00, 8_00)]
             __lowerCamelCase   =						np.asarray(UpperCamelCase_   )
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       return_tensors="""np"""   ).input_features
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       return_tensors="""np"""   ).input_features
             for enc_seq_a, enc_seq_a in zip(UpperCamelCase_    ,       UpperCamelCase_   ):
                    self.assertTrue(np.allclose(UpperCamelCase_    ,       UpperCamelCase_    ,       atol=1E-3   )   )
      def 			lowerCAmelCase__						(				self: Union[str, Any]   ):
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]

             __lowerCamelCase   =						["""longest""", """max_length""", """do_not_pad"""]
             __lowerCamelCase   =						[None, 16, None]
             for max_length, padding in zip(UpperCamelCase_    ,       UpperCamelCase_   ):
                    __lowerCamelCase   =						feature_extractor(
                        UpperCamelCase_    ,       padding=UpperCamelCase_    ,       max_length=UpperCamelCase_    ,       return_attention_mask=UpperCamelCase_   )
                    __lowerCamelCase   =						inputs.input_features
                    __lowerCamelCase   =						inputs.attention_mask
                    __lowerCamelCase   =						[np.sum(UpperCamelCase_   ) for x in attention_mask]

                    self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
                    self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
                    self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )
      def 			lowerCAmelCase__						(				self: Dict   ):
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]

             __lowerCamelCase   =						["""longest""", """max_length""", """do_not_pad"""]
             __lowerCamelCase   =						[None, 16, None]
             for max_length, padding in zip(UpperCamelCase_    ,       UpperCamelCase_   ):
                    __lowerCamelCase   =						feature_extractor(
                        UpperCamelCase_    ,       max_length=UpperCamelCase_    ,       padding=UpperCamelCase_    ,       return_tensors="""np"""    ,       return_attention_mask=UpperCamelCase_   )
                    __lowerCamelCase   =						inputs.input_features
                    __lowerCamelCase   =						inputs.attention_mask
                    __lowerCamelCase   =						[np.sum(UpperCamelCase_   ) for x in attention_mask]

                    self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]   )
                    self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6   )
                    self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]   )
                    self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6   )
                    self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]   )
      def 			lowerCAmelCase__						(				self: int   ):
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]
             __lowerCamelCase   =						feature_extractor(
                 UpperCamelCase_    ,       padding="""max_length"""    ,       max_length=4    ,       truncation=UpperCamelCase_    ,       return_tensors="""np"""    ,       return_attention_mask=UpperCamelCase_    ,       )
             __lowerCamelCase   =						inputs.input_features
             __lowerCamelCase   =						inputs.attention_mask
             __lowerCamelCase   =						np.sum(attention_mask == 1    ,       axis=1   )

             self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
             self._check_zero_mean_unit_variance(input_features[1]   )
             self._check_zero_mean_unit_variance(input_features[2]   )
      def 			lowerCAmelCase__						(				self: int   ):
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]
             __lowerCamelCase   =						feature_extractor(
                 UpperCamelCase_    ,       padding="""longest"""    ,       max_length=4    ,       truncation=UpperCamelCase_    ,       return_tensors="""np"""    ,       return_attention_mask=UpperCamelCase_    ,       )
             __lowerCamelCase   =						inputs.input_features
             __lowerCamelCase   =						inputs.attention_mask
             __lowerCamelCase   =						np.sum(attention_mask == 1    ,       axis=1   )

             self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
             self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
             self._check_zero_mean_unit_variance(input_features[2]   )

             # make sure that if max_length < longest -> then pad to max_length
             self.assertEqual(input_features.shape    ,       (3, 4, 24)   )

             __lowerCamelCase   =						[floats_list((1, x)   )[0] for x in range(8_00    ,       14_00    ,       2_00   )]
             __lowerCamelCase   =						feature_extractor(
                 UpperCamelCase_    ,       padding="""longest"""    ,       max_length=16    ,       truncation=UpperCamelCase_    ,       return_tensors="""np"""    ,       return_attention_mask=UpperCamelCase_    ,       )
             __lowerCamelCase   =						inputs.input_features
             __lowerCamelCase   =						inputs.attention_mask
             __lowerCamelCase   =						np.sum(attention_mask == 1    ,       axis=1   )

             self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]   )
             self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]   )
             self._check_zero_mean_unit_variance(input_features[2]   )

             # make sure that if max_length < longest -> then pad to max_length
             self.assertEqual(input_features.shape    ,       (3, 6, 24)   )
      def 			lowerCAmelCase__						(				self: List[Any]   ):
             import torch

             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						np.random.rand(1_00    ,       32   ).astype(np.floataa   )
             __lowerCamelCase   =						np_speech_inputs.tolist()

             for inputs in [py_speech_inputs, np_speech_inputs]:
                    __lowerCamelCase   =						feature_extractor.pad([{"""input_features""": inputs}]    ,       return_tensors="""np"""   )
                    self.assertTrue(np_processed.input_features.dtype == np.floataa   )
                    __lowerCamelCase   =						feature_extractor.pad([{"""input_features""": inputs}]    ,       return_tensors="""pt"""   )
                    self.assertTrue(pt_processed.input_features.dtype == torch.floataa   )
      def 			lowerCAmelCase__						(				self: Union[str, Any]    ,       UpperCamelCase_: Tuple   ):
             from datasets import load_dataset

             __lowerCamelCase   =						load_dataset("""hf-internal-testing/librispeech_asr_dummy"""    ,       """clean"""    ,       split="""validation"""   )
             # automatic decoding with librispeech
             __lowerCamelCase   =						ds.sort("""id"""   ).select(range(UpperCamelCase_   )   )[:num_samples]["""audio"""]

             return [x["array"] for x in speech_samples]




      def 			lowerCAmelCase__						(				self: Any   ):
             # fmt: off
             __lowerCamelCase   =						np.array([
                 -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
                 -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
                 -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
             ]   )
             # fmt: on

             __lowerCamelCase   =						self._load_datasamples(1   )
             __lowerCamelCase   =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()   )
             __lowerCamelCase   =						feature_extractor(UpperCamelCase_    ,       return_tensors="""pt"""   ).input_features
             self.assertEquals(input_features.shape    ,       (1, 5_84, 24)   )
             self.assertTrue(np.allclose(input_features[0, 0, :30]    ,       UpperCamelCase_    ,       atol=1E-4   )   )







import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('>=', FSDP_PYTORCH_VERSION):
       import torch.distributed.checkpoint as dist_cp
       from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
       from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
       from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
       from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


UpperCAmelCase_									=	get_logger(__name__)




def 					lowerCamelCase__       (						A__       :					Union[str, Any]  ,			A__       :					str  ,			A__       :					Any  ,			A__       :					Dict  ,			A__       :					Any=0	):





       '''simple docstring'''



       os.makedirs(A__  ,			exist_ok=A__	)
       with FSDP.state_dict_type(
           A__  ,			fsdp_plugin.state_dict_type  ,			fsdp_plugin.state_dict_config  ,			fsdp_plugin.optim_state_dict_config	):
              __lowerCamelCase   =						model.state_dict()
              if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
                     __lowerCamelCase   =						f'{MODEL_NAME}.bin' if model_index == 0 else f'{MODEL_NAME}_{model_index}.bin'
                     __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                     if accelerator.process_index == 0:
                            logger.info(f'Saving model to {output_model_file}'	)
                            torch.save(A__  ,			A__	)
                            logger.info(f'Model saved to {output_model_file}'	)
              elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
                     __lowerCamelCase   =						(
                         f'{MODEL_NAME}_rank{accelerator.process_index}.bin'
                         if model_index == 0
                         else f'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
                     )
                     __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                     logger.info(f'Saving model to {output_model_file}'	)
                     torch.save(A__  ,			A__	)
                     logger.info(f'Model saved to {output_model_file}'	)
              elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
                     __lowerCamelCase   =						os.path.join(A__  ,			f'{MODEL_NAME}_{model_index}'	)
                     os.makedirs(A__  ,			exist_ok=A__	)
                     logger.info(f'Saving model to {ckpt_dir}'	)
                     __lowerCamelCase   =						{"""model""": state_dict}

                     dist_cp.save_state_dict(
                         state_dict=A__  ,			storage_writer=dist_cp.FileSystemWriter(A__	)  ,			planner=DefaultSavePlanner()  ,			)
                     logger.info(f'Model saved to {ckpt_dir}'	)




def 					lowerCamelCase__       (						A__       :					int  ,			A__       :					Dict  ,			A__       :					int  ,			A__       :					List[str]  ,			A__       :					Any=0	):





       '''simple docstring'''



       accelerator.wait_for_everyone()
       with FSDP.state_dict_type(
           A__  ,			fsdp_plugin.state_dict_type  ,			fsdp_plugin.state_dict_config  ,			fsdp_plugin.optim_state_dict_config	):
              if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
                     if type(A__	) != FSDP and accelerator.process_index != 0:
                            if not fsdp_plugin.sync_module_states:
                                   raise ValueError(
                                       """Set the `sync_module_states` flag to `True` so that model states are synced across processes when """
                                       """initializing FSDP object"""	)
                            return
                     __lowerCamelCase   =						f'{MODEL_NAME}.bin' if model_index == 0 else f'{MODEL_NAME}_{model_index}.bin'
                     __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                     logger.info(f'Loading model from {input_model_file}'	)
                     __lowerCamelCase   =						torch.load(A__	)
                     logger.info(f'Model loaded from {input_model_file}'	)
              elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
                     __lowerCamelCase   =						(
                         f'{MODEL_NAME}_rank{accelerator.process_index}.bin'
                         if model_index == 0
                         else f'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
                     )
                     __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                     logger.info(f'Loading model from {input_model_file}'	)
                     __lowerCamelCase   =						torch.load(A__	)
                     logger.info(f'Model loaded from {input_model_file}'	)
              elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
                     __lowerCamelCase   =						(
                         os.path.join(A__  ,			f'{MODEL_NAME}_{model_index}'	)
                         if f'{MODEL_NAME}' not in input_dir
                         else input_dir
                     )
                     logger.info(f'Loading model from {ckpt_dir}'	)
                     __lowerCamelCase   =						{"""model""": model.state_dict()}
                     dist_cp.load_state_dict(
                         state_dict=A__  ,			storage_reader=dist_cp.FileSystemReader(A__	)  ,			planner=DefaultLoadPlanner()  ,			)
                     __lowerCamelCase   =						state_dict["""model"""]
                     logger.info(f'Model loaded from {ckpt_dir}'	)
              model.load_state_dict(A__	)




def 					lowerCamelCase__       (						A__       :					List[str]  ,			A__       :					List[str]  ,			A__       :					str  ,			A__       :					Dict  ,			A__       :					Optional[Any]  ,			A__       :					Optional[int]=0	):





       '''simple docstring'''



       os.makedirs(A__  ,			exist_ok=A__	)
       with FSDP.state_dict_type(
           A__  ,			fsdp_plugin.state_dict_type  ,			fsdp_plugin.state_dict_config  ,			fsdp_plugin.optim_state_dict_config	):
              __lowerCamelCase   =						FSDP.optim_state_dict(A__  ,			A__	)
              if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
                     if accelerator.process_index == 0:
                            __lowerCamelCase   =						(
                                f'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else f'{OPTIMIZER_NAME}_{optimizer_index}.bin'
                            )
                            __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                            logger.info(f'Saving Optimizer state to {output_optimizer_file}'	)
                            torch.save(A__  ,			A__	)
                            logger.info(f'Optimizer state saved in {output_optimizer_file}'	)
              else:
                     __lowerCamelCase   =						os.path.join(A__  ,			f'{OPTIMIZER_NAME}_{optimizer_index}'	)
                     os.makedirs(A__  ,			exist_ok=A__	)
                     logger.info(f'Saving Optimizer state to {ckpt_dir}'	)
                     dist_cp.save_state_dict(
                         state_dict={"""optimizer""": optim_state}  ,			storage_writer=dist_cp.FileSystemWriter(A__	)  ,			planner=DefaultSavePlanner()  ,			)
                     logger.info(f'Optimizer state saved in {ckpt_dir}'	)




def 					lowerCamelCase__       (						A__       :					int  ,			A__       :					List[str]  ,			A__       :					int  ,			A__       :					Any  ,			A__       :					Union[str, Any]  ,			A__       :					List[Any]=0	):





       '''simple docstring'''



       accelerator.wait_for_everyone()
       with FSDP.state_dict_type(
           A__  ,			fsdp_plugin.state_dict_type  ,			fsdp_plugin.state_dict_config  ,			fsdp_plugin.optim_state_dict_config	):
              if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
                     __lowerCamelCase   =						None
                     # below check should work but currently it isn't working (mostly opytorch issue),
                     # in the meantime disabling it at the cost of excess memory usage
                     # if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
                     __lowerCamelCase   =						(
                         f'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else f'{OPTIMIZER_NAME}_{optimizer_index}.bin'
                     )
                     __lowerCamelCase   =						os.path.join(A__  ,			A__	)
                     logger.info(f'Loading Optimizer state from {input_optimizer_file}'	)
                     __lowerCamelCase   =						torch.load(A__	)
                     logger.info(f'Optimizer state loaded from {input_optimizer_file}'	)
              else:
                     __lowerCamelCase   =						(
                         os.path.join(A__  ,			f'{OPTIMIZER_NAME}_{optimizer_index}'	)
                         if f'{OPTIMIZER_NAME}' not in input_dir
                         else input_dir
                     )
                     logger.info(f'Loading Optimizer from {ckpt_dir}'	)
                     __lowerCamelCase   =						load_sharded_optimizer_state_dict(
                         model_state_dict=model.state_dict()  ,			optimizer_key="""optimizer"""  ,			storage_reader=dist_cp.FileSystemReader(A__	)  ,			)
                     __lowerCamelCase   =						optim_state["""optimizer"""]
                     logger.info(f'Optimizer loaded from {ckpt_dir}'	)
              __lowerCamelCase   =						FSDP.optim_state_dict_to_load(A__  ,			A__  ,			A__	)
              optimizer.load_state_dict(A__	)

