



import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_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_speech_available():
 from transformers import WhisperFeatureExtractor

if is_torch_available():
 import torch

SCREAMING_SNAKE_CASE__      					=     random.Random()




def        lowerCamelCase  (				_snake_case  :						int     ,_snake_case  :						Optional[Any]=1.0     ,_snake_case  :						Optional[int]=None     ,_snake_case  :						Optional[Any]=None       ):





    '''simple docstring'''


    if rng is None:
        lowercase__		       =						global_rng

    lowercase__		       =						[]
    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 				snake_case (unittest.TestCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_=7     ,UpperCAmelCase_=400     ,UpperCAmelCase_=2_000     ,UpperCAmelCase_=10     ,UpperCAmelCase_=160     ,UpperCAmelCase_=8     ,UpperCAmelCase_=0.0     ,UpperCAmelCase_=4_000     ,UpperCAmelCase_=False     ,UpperCAmelCase_=True     ,)  ->							List[Any]:
           lowercase__		       =						parent
           lowercase__		       =						batch_size
           lowercase__		       =						min_seq_length
           lowercase__		       =						max_seq_length
           lowercase__		       =						(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
           lowercase__		       =						padding_value
           lowercase__		       =						sampling_rate
           lowercase__		       =						return_attention_mask
           lowercase__		       =						do_normalize
           lowercase__		       =						feature_size
           lowercase__		       =						chunk_length
           lowercase__		       =						hop_length



       def 							_a      (				self      )  ->							Optional[int]:
           return {
               "feature_size": self.feature_size,
               "hop_length": self.hop_length,
               "chunk_length": self.chunk_length,
               "padding_value": self.padding_value,
               "sampling_rate": self.sampling_rate,
               "return_attention_mask": self.return_attention_mask,
               "do_normalize": self.do_normalize,
           }







       def 							_a      (				self     ,UpperCAmelCase_=False     ,UpperCAmelCase_=False      )  ->							Union[str, Any]:


           def _flatten(UpperCAmelCase_      ):
               return list(itertools.chain(*_UpperCamelCase      )      )

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





@require_torch
@require_torchaudio
class 				snake_case (__lowerCAmelCase		,      unittest.TestCase   ):
       lowerCAmelCase__  :str						   =					WhisperFeatureExtractor if is_speech_available() else None



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						WhisperFeatureExtractionTester(self      )



       def 							_a      (				self      )  ->							Optional[Any]:
           lowercase__		       =						self.feature_extraction_class(**self.feat_extract_dict      )

           with tempfile.TemporaryDirectory() as tmpdirname:
               lowercase__		       =						feat_extract_first.save_pretrained(_UpperCamelCase      )[0]
               check_json_file_has_correct_format(_UpperCamelCase      )
               lowercase__		       =						self.feature_extraction_class.from_pretrained(_UpperCamelCase      )

           lowercase__		       =						feat_extract_first.to_dict()
           lowercase__		       =						feat_extract_second.to_dict()
           lowercase__		       =						feat_extract_first.mel_filters
           lowercase__		       =						feat_extract_second.mel_filters
           self.assertTrue(np.allclose(_UpperCamelCase     ,_UpperCamelCase      )      )
           self.assertEqual(_UpperCamelCase     ,_UpperCamelCase      )



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						self.feature_extraction_class(**self.feat_extract_dict      )

           with tempfile.TemporaryDirectory() as tmpdirname:
               lowercase__		       =						os.path.join(_UpperCamelCase     ,"feat_extract.json"      )
               feat_extract_first.to_json_file(_UpperCamelCase      )
               lowercase__		       =						self.feature_extraction_class.from_json_file(_UpperCamelCase      )

           lowercase__		       =						feat_extract_first.to_dict()
           lowercase__		       =						feat_extract_second.to_dict()
           lowercase__		       =						feat_extract_first.mel_filters
           lowercase__		       =						feat_extract_second.mel_filters
           self.assertTrue(np.allclose(_UpperCamelCase     ,_UpperCamelCase      )      )
           self.assertEqual(_UpperCamelCase     ,_UpperCamelCase      )



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()      )
           # create three inputs of length 800, 1000, and 1200
           lowercase__		       =						[floats_list((1, x)      )[0] for x in range(800     ,1_400     ,200      )]
           lowercase__		       =						[np.asarray(_UpperCamelCase      ) for speech_input in speech_inputs]

           # Test feature size
           lowercase__		       =						feature_extractor(_UpperCamelCase     ,padding="max_length"     ,return_tensors="np"      ).input_features
           self.assertTrue(input_features.ndim == 3      )
           self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames      )
           self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size      )

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

           # Test batched
           lowercase__		       =						feature_extractor(_UpperCamelCase     ,return_tensors="np"      ).input_features
           lowercase__		       =						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.
           lowercase__		       =						[floats_list((1, x)      )[0] for x in (800, 800, 800)]
           lowercase__		       =						np.asarray(_UpperCamelCase      )
           lowercase__		       =						feature_extractor(_UpperCamelCase     ,return_tensors="np"      ).input_features
           lowercase__		       =						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 truncation required
           lowercase__		       =						[floats_list((1, x)      )[0] for x in range(200     ,(feature_extractor.n_samples + 500)     ,200      )]
           lowercase__		       =						[np.asarray(_UpperCamelCase      ) for speech_input in speech_inputs]

           lowercase__		       =						[x[: feature_extractor.n_samples] for x in speech_inputs]
           lowercase__		       =						[np.asarray(_UpperCamelCase      ) for speech_input in speech_inputs_truncated]

           lowercase__		       =						feature_extractor(_UpperCamelCase     ,return_tensors="np"      ).input_features
           lowercase__		       =						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 							_a      (				self      )  ->							Tuple:
           import torch

           lowercase__		       =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()      )
           lowercase__		       =						np.random.rand(100     ,32      ).astype(np.floataa      )
           lowercase__		       =						np_speech_inputs.tolist()

           for inputs in [py_speech_inputs, np_speech_inputs]:
               lowercase__		       =						feature_extractor.pad([{"input_features": inputs}]     ,return_tensors="np"      )
               self.assertTrue(np_processed.input_features.dtype == np.floataa      )
               lowercase__		       =						feature_extractor.pad([{"input_features": inputs}]     ,return_tensors="pt"      )
               self.assertTrue(pt_processed.input_features.dtype == torch.floataa      )



       def 							_a      (				self     ,UpperCAmelCase_      )  ->							Union[str, Any]:
           lowercase__		       =						load_dataset("hf-internal-testing/librispeech_asr_dummy"     ,"clean"     ,split="validation"      )
           # automatic decoding with librispeech
           lowercase__		       =						ds.sort("id"      ).select(range(_UpperCamelCase      )      )[:num_samples]["""audio"""]

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



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						torch.tensor(
               [
                   0.11_93, -0.09_46, -0.10_98, -0.01_96, 0.02_25, -0.06_90, -0.17_36, 0.09_51,
                   0.09_71, -0.08_17, -0.07_02, 0.01_62, 0.02_60, 0.00_17, -0.01_92, -0.16_78,
                   0.07_09, -0.18_67, -0.06_55, -0.02_74, -0.02_34, -0.18_84, -0.05_16, -0.05_54,
                   -0.02_74, -0.14_25, -0.14_23, 0.08_37, 0.03_77, -0.08_54
               ]      )
           # fmt: on

           lowercase__		       =						self._load_datasamples(1      )
           lowercase__		       =						WhisperFeatureExtractor()
           lowercase__		       =						feature_extractor(_UpperCamelCase     ,return_tensors="pt"      ).input_features
           self.assertEqual(input_features.shape     ,(1, 80, 3_000)      )
           self.assertTrue(torch.allclose(input_features[0, 0, :30]     ,_UpperCamelCase     ,atol=1E-4      )      )







       def 							_a      (				self      )  ->							Union[str, Any]:
           lowercase__		       =						self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()      )
           lowercase__		       =						self._load_datasamples(1      )[0]
           lowercase__		       =						((audio - audio.min()) / (audio.max() - audio.min())) * 65_535  # Rescale to [0, 65535] to show issue
           lowercase__		       =						feat_extract.zero_mean_unit_var_norm([audio]     ,attention_mask=_UpperCamelCase      )[0]

           self.assertTrue(np.all(np.mean(_UpperCamelCase      ) < 1E-3      )      )
           self.assertTrue(np.all(np.abs(np.var(_UpperCamelCase      ) - 1      ) < 1E-3      )      )









'''simple docstring'''





import numpy as np
import torch
from torch.utils.data import Dataset

from utils import logger


class 				snake_case (UpperCamelCase   ):



       def __init__(				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							int:
           lowercase__		       =						params

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array([len(UpperCAmelCase_      ) for t in data]      )

           self.check()
           self.remove_long_sequences()
           self.remove_empty_sequences()
           self.remove_unknown_sequences()
           self.check()
           self.print_statistics()



       def __getitem__(				self     ,UpperCAmelCase_      )  ->							Optional[int]:
           return (self.token_ids[index], self.lengths[index])



       def __len__(				self      )  ->							Union[str, Any]:
           return len(self.lengths      )



       def 							_a      (				self      )  ->							int:
           assert len(self.token_ids      ) == len(self.lengths      )
           assert all(self.lengths[i] == len(self.token_ids[i]      ) for i in range(len(self.lengths      )      )      )



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						self.params.max_model_input_size
           lowercase__		       =						self.lengths > max_len
           logger.info(F'''Splitting {sum(UpperCAmelCase_      )} too long sequences.'''      )

           def divide_chunks(UpperCAmelCase_     ,UpperCAmelCase_      ):
               return [l[i : i + n] for i in range(0     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )]

           lowercase__		       =						[]
           lowercase__		       =						[]
           if self.params.mlm:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["cls_token"], self.params.special_tok_ids["sep_token"]
           else:
               lowercase__       ,       lowercase__		       =						self.params.special_tok_ids["bos_token"], self.params.special_tok_ids["eos_token"]

           for seq_, len_ in zip(self.token_ids     ,self.lengths      ):
               assert (seq_[0] == cls_id) and (seq_[-1] == sep_id), seq_
               if len_ <= max_len:
                   new_tok_ids.append(seq_      )
                   new_lengths.append(len_      )
               else:
                   lowercase__		       =						[]
                   for sub_s in divide_chunks(seq_     ,max_len - 2      ):
                       if sub_s[0] != cls_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,0     ,UpperCAmelCase_      )
                       if sub_s[-1] != sep_id:
                           lowercase__		       =						np.insert(UpperCAmelCase_     ,len(UpperCAmelCase_      )     ,UpperCAmelCase_      )
                       assert len(UpperCAmelCase_      ) <= max_len
                       assert (sub_s[0] == cls_id) and (sub_s[-1] == sep_id), sub_s
                       sub_seqs.append(UpperCAmelCase_      )

                   new_tok_ids.extend(UpperCAmelCase_      )
                   new_lengths.extend([len(UpperCAmelCase_      ) for l in sub_seqs]      )

           lowercase__		       =						np.array(UpperCAmelCase_      )
           lowercase__		       =						np.array(UpperCAmelCase_      )



       def 							_a      (				self      )  ->							Any:
           lowercase__		       =						len(self      )
           lowercase__		       =						self.lengths > 11
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} too short (<=11 tokens) sequences.'''      )



       def 							_a      (				self      )  ->							List[Any]:
           if "unk_token" not in self.params.special_tok_ids:
               return
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						len(self      )
           lowercase__		       =						np.array([np.count_nonzero(a == unk_token_id      ) for a in self.token_ids]      )
           lowercase__		       =						(unk_occs / self.lengths) < 0.5
           lowercase__		       =						self.token_ids[indices]
           lowercase__		       =						self.lengths[indices]
           lowercase__		       =						len(self      )
           logger.info(F'''Remove {init_size - new_size} sequences with a high level of unknown tokens (50%).'''      )



       def 							_a      (				self      )  ->							Optional[int]:
           if not self.params.is_master:
               return
           logger.info(F'''{len(self      )} sequences'''      )
           # data_len = sum(self.lengths)
           # nb_unique_tokens = len(Counter(list(chain(*self.token_ids))))
           # logger.info(f'{data_len} tokens ({nb_unique_tokens} unique)')

           # unk_idx = self.params.special_tok_ids['unk_token']
           # nb_unknown = sum([(t==unk_idx).sum() for t in self.token_ids])
           # logger.info(f'{nb_unknown} unknown tokens (covering {100*nb_unknown/data_len:.2f}% of the data)')







       def 							_a      (				self     ,UpperCAmelCase_      )  ->							List[str]:
           lowercase__		       =						[t[0] for t in batch]
           lowercase__		       =						[t[1] for t in batch]
           assert len(UpperCAmelCase_      ) == len(UpperCAmelCase_      )

           # Max for paddings
           lowercase__		       =						max(UpperCAmelCase_      )

           # Pad token ids
           if self.params.mlm:
               lowercase__		       =						self.params.special_tok_ids["pad_token"]
           else:
               lowercase__		       =						self.params.special_tok_ids["unk_token"]
           lowercase__		       =						[list(t.astype(UpperCAmelCase_      )      ) + [pad_idx] * (max_seq_len_ - len(UpperCAmelCase_      )) for t in token_ids]
           assert len(tk_      ) == len(UpperCAmelCase_      )
           assert all(len(UpperCAmelCase_      ) == max_seq_len_ for t in tk_      )

           lowercase__		       =						torch.tensor(tk_      )  # (bs, max_seq_len_)
           lowercase__		       =						torch.tensor(UpperCAmelCase_      )  # (bs)
           return tk_t, lg_t




