

"""simple docstring"""

# XXX: we want transformers master here - in the absense of conftest manipulating sys.path:
# hack it in for now:
import sys
from pathlib import Path


_UpperCamelCase		     =    Path(__file__).resolve().parents[3] / """src"""
sys.path.insert(1, str(git_repo_path))

import dataclasses  # noqa
import io  # noqa
import itertools  # noqa
import json  # noqa
import os  # noqa
import unittest  # noqa
from copy import deepcopy  # noqa

from parameterized import parameterized  # noqa
from transformers import TrainingArguments, is_torch_available  # noqa
from transformers.deepspeed import is_deepspeed_available  # noqa
from transformers.file_utils import WEIGHTS_NAME  # noqa
from transformers.testing_utils import (  # noqa
    CaptureLogger,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    mockenv_context,
    require_deepspeed,
    require_torch_gpu,
    require_torch_multi_gpu,
    slow,
)
from transformers.trainer_utils import set_seed  # noqa


set_seed(4_2)

_UpperCamelCase		     =    {"""base""": """patrickvonplaten/wav2vec2_tiny_random""", """robust""": """patrickvonplaten/wav2vec2_tiny_random_robust"""}

_UpperCamelCase		     =    """zero2"""
_UpperCamelCase		     =    """zero3"""
_UpperCamelCase		     =    [ZEROa, ZEROa]



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__					,					lowercase__   )			->						Dict:
 # customize the test name generator function as we want both params to appear in the sub-test
 # name, as by default it shows only the first param
 lowerCAmelCase__   :				Dict 			=				parameterized.to_safe_name("_".join(str(lowercase__   ) for x in param.args   )   )
 return F"""{func.__name__}_{param_based_name}"""


# Cartesian-product of zero stages with models to test
_UpperCamelCase		     =    list(itertools.product(stages, models.keys()))






@slow
@require_deepspeed
@require_torch_gpu
class 			__a					(						__magic_name__       ):





       """simple docstring"""


       @parameterized.expand(snake_case			,					name_func=snake_case   )
       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case   ):





        """simple docstring"""

        self.run_and_check(
            stage=snake_case			,					model=snake_case			,					distributed=snake_case			,					fpaa=snake_case			,					)






       @require_torch_multi_gpu
       @parameterized.expand(snake_case			,					name_func=snake_case   )
       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case   ):





        """simple docstring"""

        self.run_and_check(
            stage=snake_case			,					model=snake_case			,					distributed=snake_case			,					fpaa=snake_case			,					)






       @parameterized.expand(snake_case			,					name_func=snake_case   )
       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case   ):





        """simple docstring"""

        self.run_and_check(
            stage=snake_case			,					model=snake_case			,					distributed=snake_case			,					fpaa=snake_case			,					)






       @require_torch_multi_gpu
       @parameterized.expand(snake_case			,					name_func=snake_case   )
       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case   ):





        """simple docstring"""

        self.run_and_check(
            stage=snake_case			,					model=snake_case			,					distributed=snake_case			,					fpaa=snake_case			,					)






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

        pass






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case			,					snake_case = 10			,					snake_case = True			,					snake_case = True			,					snake_case = True			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Dict 			=				models[model]

        lowerCAmelCase__   :				Dict 			=				self.run_trainer(
            stage=snake_case			,					model_name=snake_case			,					eval_steps=snake_case			,					num_train_epochs=1			,					distributed=snake_case			,					fpaa=snake_case			,					)

        self.do_checks(snake_case   )

        return output_dir






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case			,					snake_case = 10			,					snake_case = 1			,					snake_case = True			,					snake_case = True			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Any 			=				self.get_auto_remove_tmp_dir("./xxx"			,					after=snake_case   )
        lowerCAmelCase__   :				List[str] 			=				F"""
            --model_name_or_path {model_name}
            --dataset_name hf-internal-testing/librispeech_asr_dummy
            --dataset_config_name clean
            --train_split_name validation
            --validation_split_name validation
            --output_dir {output_dir}
            --num_train_epochs {str(snake_case   )}
            --per_device_train_batch_size 2
            --per_device_eval_batch_size 2
            --evaluation_strategy steps
            --learning_rate 5e-4
            --warmup_steps 8
            --orthography timit
            --preprocessing_num_workers 1
            --group_by_length
            --freeze_feature_extractor
            --report_to none
            --save_steps 0
            --eval_steps {eval_steps}
            --report_to none
        """.split()

        if fpaa:
         args.extend(["--fp16"]   )

        # currently ds_config_wav2vec2_zero.json requires "zero_optimization.find_unused_parameters": true,
        # hence the separate config files
        lowerCAmelCase__   :				Dict 			=				F"""--deepspeed {self.test_file_dir_str}/ds_config_wav2vec2_{stage}.json""".split()
        lowerCAmelCase__   :				str 			=				[F"""{self.examples_dir_str}/research_projects/wav2vec2/run_asr.py"""]
        lowerCAmelCase__   :				Optional[Any] 			=				self.get_launcher(snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				launcher + script + args + ds_args
        # keep for quick debug
        # print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
        execute_subprocess_async(snake_case			,					env=self.get_env()   )

        return output_dir






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case=False   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				min(2			,					get_gpu_count()   ) if distributed else 1
        return F"""deepspeed --num_nodes 1 --num_gpus {num_gpus}""".split()









"""simple docstring"""

import itertools
import os
import random
import tempfile
import unittest

import numpy as np

from transformers import TvltFeatureExtractor, is_datasets_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
    import torch

if is_datasets_available():
    from datasets import load_dataset

_UpperCamelCase		     =    random.Random()



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__=1.0					,					lowercase__=None					,					lowercase__=None   )			->						Optional[Any]:
 if rng is None:
  lowerCAmelCase__   :				Union[str, Any] 			=				global_rng

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

 return values






class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def __init__(   self			,					snake_case			,					snake_case=7			,					snake_case=400			,					snake_case=2_000			,					snake_case=2_048			,					snake_case=128			,					snake_case=1			,					snake_case=512			,					snake_case=30			,					snake_case=44_100			,					):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				parent
        lowerCAmelCase__   :				Optional[int] 			=				batch_size
        lowerCAmelCase__   :				Optional[Any] 			=				min_seq_length
        lowerCAmelCase__   :				Optional[int] 			=				max_seq_length
        lowerCAmelCase__   :				Dict 			=				(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
        lowerCAmelCase__   :				Dict 			=				spectrogram_length
        lowerCAmelCase__   :				Any 			=				feature_size
        lowerCAmelCase__   :				int 			=				num_audio_channels
        lowerCAmelCase__   :				Optional[int] 			=				hop_length
        lowerCAmelCase__   :				List[str] 			=				chunk_length
        lowerCAmelCase__   :				Optional[Any] 			=				sampling_rate






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return {
            "spectrogram_length": self.spectrogram_length,
            "feature_size": self.feature_size,
            "num_audio_channels": self.num_audio_channels,
            "hop_length": self.hop_length,
            "chunk_length": self.chunk_length,
            "sampling_rate": self.sampling_rate,
        }






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case=False			,					snake_case=False   ):





        """simple docstring"""

        def _flatten(snake_case   ):
         return list(itertools.chain(*snake_case   )   )

        if equal_length:
         lowerCAmelCase__   :				Dict 			=				[floats_list((self.max_seq_length, self.feature_size)   ) for _ in range(self.batch_size   )]
        else:
         # make sure that inputs increase in size
         lowerCAmelCase__   :				Tuple 			=				[
             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__   :				Dict 			=				[np.asarray(snake_case   ) for x in speech_inputs]
        return speech_inputs






@require_torch
@require_torchaudio
class 			__a					(						__magic_name__       ,	unittest.TestCase       ):





       """simple docstring"""


       __UpperCamelCase					:      Optional[Any]    =						TvltFeatureExtractor






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				TvltFeatureExtractionTester(self   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				self.feature_extraction_class(**self.feat_extract_dict   )
        self.assertTrue(hasattr(snake_case			,					"spectrogram_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"feature_size"   )   )
        self.assertTrue(hasattr(snake_case			,					"num_audio_channels"   )   )
        self.assertTrue(hasattr(snake_case			,					"hop_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"chunk_length"   )   )
        self.assertTrue(hasattr(snake_case			,					"sampling_rate"   )   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				int 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				Union[str, Any] 			=				feat_extract_first.save_pretrained(snake_case   )[0]
         check_json_file_has_correct_format(snake_case   )
         lowerCAmelCase__   :				Tuple 			=				self.feature_extraction_class.from_pretrained(snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				feat_extract_first.to_dict()
        lowerCAmelCase__   :				int 			=				feat_extract_second.to_dict()
        lowerCAmelCase__   :				Union[str, Any] 			=				dict_first.pop("mel_filters"   )
        lowerCAmelCase__   :				List[Any] 			=				dict_second.pop("mel_filters"   )
        self.assertTrue(np.allclose(snake_case			,					snake_case   )   )
        self.assertEqual(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Any 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        with tempfile.TemporaryDirectory() as tmpdirname:
         lowerCAmelCase__   :				List[Any] 			=				os.path.join(snake_case			,					"feat_extract.json"   )
         feat_extract_first.to_json_file(snake_case   )
         lowerCAmelCase__   :				List[Any] 			=				self.feature_extraction_class.from_json_file(snake_case   )

        lowerCAmelCase__   :				Optional[Any] 			=				feat_extract_first.to_dict()
        lowerCAmelCase__   :				Dict 			=				feat_extract_second.to_dict()
        lowerCAmelCase__   :				Tuple 			=				dict_first.pop("mel_filters"   )
        lowerCAmelCase__   :				Dict 			=				dict_second.pop("mel_filters"   )
        self.assertTrue(np.allclose(snake_case			,					snake_case   )   )
        self.assertEqual(snake_case			,					snake_case   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.feature_extraction_class(**self.feat_extract_dict   )

        # create three inputs of length 800, 1000, and 1200
        lowerCAmelCase__   :				Union[str, Any] 			=				[floats_list((1, x)   )[0] for x in range(800			,					1_400			,					200   )]
        lowerCAmelCase__   :				Optional[int] 			=				[np.asarray(snake_case   ) for speech_input in speech_inputs]

        # Test not batched input
        lowerCAmelCase__   :				Tuple 			=				feature_extractor(np_speech_inputs[0]			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test batched
        lowerCAmelCase__   :				Optional[int] 			=				feature_extractor(snake_case			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test audio masking
        lowerCAmelCase__   :				List[str] 			=				feature_extractor(
            snake_case			,					return_tensors="np"			,					sampling_rate=44_100			,					mask_audio=snake_case   ).audio_values

        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )

        # Test 2-D numpy arrays are batched.
        lowerCAmelCase__   :				List[str] 			=				[floats_list((1, x)   )[0] for x in (800, 800, 800)]
        lowerCAmelCase__   :				Tuple 			=				np.asarray(snake_case   )
        lowerCAmelCase__   :				List[Any] 			=				feature_extractor(snake_case			,					return_tensors="np"			,					sampling_rate=44_100   ).audio_values
        self.assertTrue(encoded_audios.ndim == 4   )
        self.assertTrue(encoded_audios.shape[-1] == feature_extractor.feature_size   )
        self.assertTrue(encoded_audios.shape[-2] <= feature_extractor.spectrogram_length   )
        self.assertTrue(encoded_audios.shape[-3] == feature_extractor.num_channels   )






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case   ):





        """simple docstring"""

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

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






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				str 			=				self._load_datasamples(1   )
        lowerCAmelCase__   :				Optional[Any] 			=				TvltFeatureExtractor()
        lowerCAmelCase__   :				Dict 			=				feature_extractor(snake_case			,					return_tensors="pt"   ).audio_values

        self.assertEquals(audio_values.shape			,					(1, 1, 192, 128)   )

        lowerCAmelCase__   :				Dict 			=				torch.tensor([[-0.3_032, -0.2_708], [-0.4_434, -0.4_007]]   )
        self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2]			,					snake_case			,					atol=1e-4   )   )






