




"""simple docstring"""





import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
							from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


_lowercase	      =    get_tests_dir("""fixtures/test_sentencepiece.model""")

@require_torch
@require_torchaudio
@require_sentencepiece
class       __lowercase       (			unittest.TestCase	):


				'''simple docstring'''


				def       _lowerCamelCase					(   self   :     List[Any]      ):
										snake_case__       :					str										=					tempfile.mkdtemp()

										snake_case__       :					List[Any]										=					['''<s>''', '''<pad>''', '''</s>''', '''<unk>''', '''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''']
										snake_case__       :					List[str]										=					dict(zip(SCREAMING_SNAKE_CASE_      ,range(len(SCREAMING_SNAKE_CASE_      )      )      )      )
										snake_case__       :					Dict										=					Path(self.tmpdirname      )
										save_json(SCREAMING_SNAKE_CASE_      ,save_dir / VOCAB_FILES_NAMES['''vocab_file''']      )
										if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
																copyfile(SCREAMING_SNAKE_CASE_      ,save_dir / VOCAB_FILES_NAMES['''spm_file''']      )

										snake_case__       :					List[str]										=					SpeechaTextTokenizer.from_pretrained(self.tmpdirname      )
										tokenizer.save_pretrained(self.tmpdirname      )

										snake_case__       :					int										=					{
										    '''feature_size''': 2_4,
										    '''num_mel_bins''': 2_4,
										    '''padding_value''': 0.0,
										    '''sampling_rate''': 1_6_0_0_0,
										    '''return_attention_mask''': False,
										    '''do_normalize''': True,
										}
										save_json(SCREAMING_SNAKE_CASE_      ,save_dir / FEATURE_EXTRACTOR_NAME      )

				def       _lowerCamelCase					(   self   :     Union[str, Any]      ,**SCREAMING_SNAKE_CASE_   :     Union[str, Any]      ):
										return SpeechaTextTokenizer.from_pretrained(self.tmpdirname      ,**SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     int      ,**SCREAMING_SNAKE_CASE_   :     Dict      ):
										return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname      ,**SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Optional[int]      ):
										shutil.rmtree(self.tmpdirname      )

				def       _lowerCamelCase					(   self   :     Optional[Any]      ):
										snake_case__       :					List[str]										=					self.get_tokenizer()
										snake_case__       :					List[str]										=					self.get_feature_extractor()

										snake_case__       :					Optional[Any]										=					SpeechaTextProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										processor.save_pretrained(self.tmpdirname      )
										snake_case__       :					Optional[Any]										=					SpeechaTextProcessor.from_pretrained(self.tmpdirname      )

										self.assertEqual(processor.tokenizer.get_vocab()      ,tokenizer.get_vocab()      )
										self.assertIsInstance(processor.tokenizer      ,SCREAMING_SNAKE_CASE_      )

										self.assertEqual(processor.feature_extractor.to_json_string()      ,feature_extractor.to_json_string()      )
										self.assertIsInstance(processor.feature_extractor      ,SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										snake_case__       :					Tuple										=					SpeechaTextProcessor(
										    tokenizer=self.get_tokenizer()      ,feature_extractor=self.get_feature_extractor()      )
										processor.save_pretrained(self.tmpdirname      )

										snake_case__       :					Any										=					self.get_tokenizer(bos_token='''(BOS)'''      ,eos_token='''(EOS)'''      )
										snake_case__       :					Optional[int]										=					self.get_feature_extractor(do_normalize=SCREAMING_SNAKE_CASE_      ,padding_value=1.0      )

										snake_case__       :					List[str]										=					SpeechaTextProcessor.from_pretrained(
										    self.tmpdirname      ,bos_token='''(BOS)'''      ,eos_token='''(EOS)'''      ,do_normalize=SCREAMING_SNAKE_CASE_      ,padding_value=1.0      )

										self.assertEqual(processor.tokenizer.get_vocab()      ,tokenizer_add_kwargs.get_vocab()      )
										self.assertIsInstance(processor.tokenizer      ,SCREAMING_SNAKE_CASE_      )

										self.assertEqual(processor.feature_extractor.to_json_string()      ,feature_extractor_add_kwargs.to_json_string()      )
										self.assertIsInstance(processor.feature_extractor      ,SCREAMING_SNAKE_CASE_      )

				def       _lowerCamelCase					(   self   :     Union[str, Any]      ):
										snake_case__       :					Any										=					self.get_feature_extractor()
										snake_case__       :					str										=					self.get_tokenizer()

										snake_case__       :					Optional[int]										=					SpeechaTextProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					List[Any]										=					floats_list((3, 1_0_0_0)      )

										snake_case__       :					Optional[int]										=					feature_extractor(SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''      )
										snake_case__       :					Optional[int]										=					processor(SCREAMING_SNAKE_CASE_      ,return_tensors='''np'''      )

										for key in input_feat_extract.keys():
																self.assertAlmostEqual(input_feat_extract[key].sum()      ,input_processor[key].sum()      ,delta=1E-2      )

				def       _lowerCamelCase					(   self   :     Any      ):
										snake_case__       :					Optional[Any]										=					self.get_feature_extractor()
										snake_case__       :					List[Any]										=					self.get_tokenizer()

										snake_case__       :					Any										=					SpeechaTextProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					Union[str, Any]										=					'''This is a test string'''

										snake_case__       :					Optional[int]										=					processor(text=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					int										=					tokenizer(SCREAMING_SNAKE_CASE_      )

										for key in encoded_tok.keys():
																self.assertListEqual(encoded_tok[key]      ,encoded_processor[key]      )

				def       _lowerCamelCase					(   self   :     str      ):
										snake_case__       :					int										=					self.get_feature_extractor()
										snake_case__       :					List[str]										=					self.get_tokenizer()

										snake_case__       :					Optional[int]										=					SpeechaTextProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										snake_case__       :					List[Any]										=					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

										snake_case__       :					Optional[Any]										=					processor.batch_decode(SCREAMING_SNAKE_CASE_      )
										snake_case__       :					Dict										=					tokenizer.batch_decode(SCREAMING_SNAKE_CASE_      )

										self.assertListEqual(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      )





				def       _lowerCamelCase					(   self   :     Tuple      ):
										snake_case__       :					Union[str, Any]										=					self.get_feature_extractor()
										snake_case__       :					Any										=					self.get_tokenizer()

										snake_case__       :					int										=					SpeechaTextProcessor(tokenizer=SCREAMING_SNAKE_CASE_      ,feature_extractor=SCREAMING_SNAKE_CASE_      )

										self.assertListEqual(
										    processor.model_input_names      ,feature_extractor.model_input_names      ,msg='''`processor` and `feature_extractor` model input names do not match'''      ,)








"""simple docstring"""





from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


_lowercase	      =    logging.get_logger(__name__)

class       __lowercase       (			snake_case_	):


				'''simple docstring'''


				_lowerCAmelCase				=["input_values", "attention_mask"]

				def __init__(   self   :     Tuple      ,SCREAMING_SNAKE_CASE_   :     Dict=1      ,SCREAMING_SNAKE_CASE_   :     List[str]=1_6_0_0_0      ,SCREAMING_SNAKE_CASE_   :     Any=1_2_8      ,SCREAMING_SNAKE_CASE_   :     List[str]=1_0_2_4      ,SCREAMING_SNAKE_CASE_   :     Any=0.0      ,SCREAMING_SNAKE_CASE_   :     Optional[Any]=True      ,SCREAMING_SNAKE_CASE_   :     Tuple=-4.2677393      ,SCREAMING_SNAKE_CASE_   :     List[Any]=4.5689974      ,SCREAMING_SNAKE_CASE_   :     str=False      ,**SCREAMING_SNAKE_CASE_   :     Dict      ,):
										super().__init__(feature_size=SCREAMING_SNAKE_CASE_      ,sampling_rate=SCREAMING_SNAKE_CASE_      ,padding_value=SCREAMING_SNAKE_CASE_      ,**SCREAMING_SNAKE_CASE_      )
										snake_case__       :					str										=					num_mel_bins
										snake_case__       :					Union[str, Any]										=					max_length
										snake_case__       :					List[str]										=					do_normalize
										snake_case__       :					Optional[int]										=					mean
										snake_case__       :					int										=					std
										snake_case__       :					Any										=					return_attention_mask

				def       _lowerCamelCase					(   self   :     Any      ,SCREAMING_SNAKE_CASE_   :     np.ndarray      ,SCREAMING_SNAKE_CASE_   :     int      ,):
										snake_case__       :					Any										=					torch.from_numpy(SCREAMING_SNAKE_CASE_      ).unsqueeze(0      )
										snake_case__       :					Optional[int]										=					ta_kaldi.fbank(
										    SCREAMING_SNAKE_CASE_      ,htk_compat=SCREAMING_SNAKE_CASE_      ,sample_frequency=self.sampling_rate      ,use_energy=SCREAMING_SNAKE_CASE_      ,window_type='''hanning'''      ,num_mel_bins=self.num_mel_bins      ,dither=0.0      ,frame_shift=1_0      ,)

										snake_case__       :					Dict										=					fbank.shape[0]
										snake_case__       :					Optional[int]										=					max_length - n_frames

										# pad or truncate, depending on difference
										if difference > 0:
																snake_case__       :					List[str]										=					torch.nn.ZeroPadad((0, 0, 0, difference)      )
																snake_case__       :					List[Any]										=					pad_module(SCREAMING_SNAKE_CASE_      )
										elif difference < 0:
																snake_case__       :					List[str]										=					fbank[0:max_length, :]

										snake_case__       :					Dict										=					fbank.numpy()

										return fbank

				def       _lowerCamelCase					(   self   :     Optional[int]      ,SCREAMING_SNAKE_CASE_   :     np.ndarray      ):
										return (input_values - (self.mean)) / (self.std * 2)





				def __call__(   self   :     Tuple      ,SCREAMING_SNAKE_CASE_   :     Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]      ,SCREAMING_SNAKE_CASE_   :     Optional[int] = None      ,SCREAMING_SNAKE_CASE_   :     Optional[Union[str, TensorType]] = None      ,**SCREAMING_SNAKE_CASE_   :     int      ,):

										if sampling_rate is not None:
																if sampling_rate != self.sampling_rate:
																						raise ValueError(
																						    F'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
																						    F''' {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with'''
																						    F''' {self.sampling_rate} and not {sampling_rate}.'''      )
										else:
																logger.warning(
																    '''It is strongly recommended to pass the `sampling_rate` argument to this function. '''
																    '''Failing to do so can result in silent errors that might be hard to debug.'''      )

										snake_case__       :					Optional[Any]										=					isinstance(SCREAMING_SNAKE_CASE_      ,np.ndarray      ) and len(raw_speech.shape      ) > 1
										if is_batched_numpy and len(raw_speech.shape      ) > 2:
																raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''      )
										snake_case__       :					Union[str, Any]										=					is_batched_numpy or (
										    isinstance(SCREAMING_SNAKE_CASE_      ,(list, tuple)      ) and (isinstance(raw_speech[0]      ,(np.ndarray, tuple, list)      ))
										)

										if is_batched:
																snake_case__       :					Optional[Any]										=					[np.asarray(SCREAMING_SNAKE_CASE_      ,dtype=np.floataa      ) for speech in raw_speech]
										elif not is_batched and not isinstance(SCREAMING_SNAKE_CASE_      ,np.ndarray      ):
																snake_case__       :					Union[str, Any]										=					np.asarray(SCREAMING_SNAKE_CASE_      ,dtype=np.floataa      )
										elif isinstance(SCREAMING_SNAKE_CASE_      ,np.ndarray      ) and raw_speech.dtype is np.dtype(np.floataa      ):
																snake_case__       :					List[str]										=					raw_speech.astype(np.floataa      )

										# always return batch
										if not is_batched:
																snake_case__       :					Tuple										=					[raw_speech]

										# extract fbank features and pad/truncate to max_length
										snake_case__       :					List[Any]										=					[self._extract_fbank_features(SCREAMING_SNAKE_CASE_      ,max_length=self.max_length      ) for waveform in raw_speech]

										# convert into BatchFeature
										snake_case__       :					str										=					BatchFeature({'''input_values''': features}      )

										# make sure list is in array format
										snake_case__       :					Union[str, Any]										=					padded_inputs.get('''input_values'''      )
										if isinstance(input_values[0]      ,SCREAMING_SNAKE_CASE_      ):
																snake_case__       :					List[str]										=					[np.asarray(SCREAMING_SNAKE_CASE_      ,dtype=np.floataa      ) for feature in input_values]

										# normalization
										if self.do_normalize:
																snake_case__       :					List[Any]										=					[self.normalize(SCREAMING_SNAKE_CASE_      ) for feature in input_values]

										if return_tensors is not None:
																snake_case__       :					Union[str, Any]										=					padded_inputs.convert_to_tensors(SCREAMING_SNAKE_CASE_      )

										return padded_inputs


