






import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__snake_case      :      Tuple             =       logging.get_logger(__name__)

__snake_case      :      Tuple             =       """▁"""

__snake_case      :      Dict             =       {
    """vocab_file""": """vocab.json""",
    """spm_file""": """sentencepiece.bpe.model""",
}

__snake_case      :      Union[str, Any]             =       {
    """vocab_file""": {
        """facebook/s2t-small-librispeech-asr""": (
            """https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/vocab.json"""
        ),
    },
    """spm_file""": {
        """facebook/s2t-small-librispeech-asr""": (
            """https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/sentencepiece.bpe.model"""
        )
    },
}

__snake_case      :      Union[str, Any]             =       {
    """facebook/s2t-small-librispeech-asr""": 10_24,
}

__snake_case      :      Optional[Any]             =       ["""pt""", """fr""", """ru""", """nl""", """ro""", """it""", """es""", """de"""]

__snake_case      :      Optional[int]             =       {"""mustc""": MUSTC_LANGS}


class       __SCREAMING_SNAKE_CASE   (				lowerCamelCase__):
      _SCREAMING_SNAKE_CASE		:    Optional[int]						=   VOCAB_FILES_NAMES
      _SCREAMING_SNAKE_CASE		:    Tuple						=   PRETRAINED_VOCAB_FILES_MAP
      _SCREAMING_SNAKE_CASE		:    Union[str, Any]						=   MAX_MODEL_INPUT_SIZES
      _SCREAMING_SNAKE_CASE		:    int						=   ["""input_ids""", """attention_mask"""]

      _SCREAMING_SNAKE_CASE		:    List[int]						=   []






      def __init__(		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase="<s>"       ,  _UpperCamelCase="</s>"       ,  _UpperCamelCase="<pad>"       ,  _UpperCamelCase="<unk>"       ,  _UpperCamelCase=False       ,  _UpperCamelCase=False       ,  _UpperCamelCase=None       ,  _UpperCamelCase=None       ,  _UpperCamelCase = None       ,  **_UpperCamelCase       ,  ):






            """simple docstring"""
            lowerCAmelCase__        =  {} if sp_model_kwargs is None else sp_model_kwargs

            super().__init__(
                bos_token=snake_case__       ,  eos_token=snake_case__       ,  unk_token=snake_case__       ,  pad_token=snake_case__       ,  do_upper_case=snake_case__       ,  do_lower_case=snake_case__       ,  tgt_lang=snake_case__       ,  lang_codes=snake_case__       ,  sp_model_kwargs=self.sp_model_kwargs       ,  **snake_case__       ,  )
            lowerCAmelCase__        =  do_upper_case
            lowerCAmelCase__        =  do_lower_case

            lowerCAmelCase__        =  load_json(snake_case__				)
            lowerCAmelCase__        =  {v: k for k, v in self.encoder.items()}
            lowerCAmelCase__        =  spm_file
            lowerCAmelCase__        =  load_spm(snake_case__       ,  self.sp_model_kwargs				)

            if lang_codes is not None:
                  lowerCAmelCase__        =  lang_codes
                  lowerCAmelCase__        =  LANGUAGES[lang_codes]
                  lowerCAmelCase__        =  [F"<lang:{lang}>" for lang in self.langs]
                  lowerCAmelCase__        =  {lang: self.sp_model.PieceToId(F"<lang:{lang}>"				) for lang in self.langs}

                  lowerCAmelCase__        =  self.lang_tokens
                  lowerCAmelCase__        =  tgt_lang if tgt_lang is not None else self.langs[0]

                  self.set_tgt_lang_special_tokens(self._tgt_lang				)
            else:
                  lowerCAmelCase__        =  {}






      @property
      def  UpperCamelCase__ (		self				):






            """simple docstring"""
            return len(self.encoder				)






      @property
      def  UpperCamelCase__ (		self				):






            """simple docstring"""
            return self._tgt_lang






      @tgt_lang.setter
      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            lowerCAmelCase__        =  new_tgt_lang
            self.set_tgt_lang_special_tokens(snake_case__				)






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            lowerCAmelCase__        =  self.lang_code_to_id[tgt_lang]
            lowerCAmelCase__        =  [lang_code_id]






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            return self.sp_model.encode(snake_case__       ,  out_type=snake_case__				)






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            return self.encoder.get(snake_case__       ,  self.encoder[self.unk_token]				)






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            return self.decoder.get(snake_case__       ,  self.unk_token				)






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase				):






            """simple docstring"""
            lowerCAmelCase__        =  []
            lowerCAmelCase__        =  ''''''
            for token in tokens:
                  # make sure that special tokens are not decoded using sentencepiece model
                  if token in self.all_special_tokens:
                        lowerCAmelCase__        =  self.sp_model.decode(snake_case__				)
                        out_string += (decoded.upper() if self.do_upper_case else decoded) + token + " "
                        lowerCAmelCase__        =  []
                  else:
                        current_sub_tokens.append(snake_case__				)
            lowerCAmelCase__        =  self.sp_model.decode(snake_case__				)
            out_string += decoded.upper() if self.do_upper_case else decoded
            return out_string.strip()






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase=None				):






            """simple docstring"""





            if token_ids_a is None:
                  return self.prefix_tokens + token_ids_a + [self.eos_token_id]
            # We don't expect to process pairs, but leave the pair logic for API consistency
            return self.prefix_tokens + token_ids_a + token_ids_a + [self.eos_token_id]






      def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase = None       ,  _UpperCamelCase = False				):






            """simple docstring"""





            if already_has_special_tokens:
                  return super().get_special_tokens_mask(
                      token_ids_a=snake_case__       ,  token_ids_a=snake_case__       ,  already_has_special_tokens=snake_case__				)

            lowerCAmelCase__        =  [1] * len(self.prefix_tokens				)
            lowerCAmelCase__        =  [1]
            if token_ids_a is None:
                  return prefix_ones + ([0] * len(snake_case__				)) + suffix_ones
            return prefix_ones + ([0] * len(snake_case__				)) + ([0] * len(snake_case__				)) + suffix_ones






      def  UpperCamelCase__ (		self				):






            """simple docstring"""
            lowerCAmelCase__        =  self.encoder.copy()
            vocab.update(self.added_tokens_encoder				)
            return vocab






      def __getstate__(		self				):






            """simple docstring"""
            lowerCAmelCase__        =  self.__dict__.copy()
            lowerCAmelCase__        =  None
            return state






      def __setstate__(		self       ,  _UpperCamelCase				):






            """simple docstring"""
            lowerCAmelCase__        =  d

            # for backward compatibility
            if not hasattr(self       ,  'sp_model_kwargs'				):
                  lowerCAmelCase__        =  {}

            lowerCAmelCase__        =  load_spm(self.spm_file       ,  self.sp_model_kwargs				)







      def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase = None				):






            """simple docstring"""
            lowerCAmelCase__        =  Path(snake_case__				)
            assert save_dir.is_dir(), F"{save_directory} should be a directory"
            lowerCAmelCase__        =  save_dir / (
                (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''vocab_file''']
            )
            lowerCAmelCase__        =  save_dir / (
                (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''spm_file''']
            )

            save_json(self.encoder       ,  snake_case__				)

            if os.path.abspath(self.spm_file				) != os.path.abspath(snake_case__				) and os.path.isfile(self.spm_file				):
                  copyfile(self.spm_file       ,  snake_case__				)
            elif not os.path.isfile(self.spm_file				):
                  with open(snake_case__       ,  'wb'				) as fi:
                        lowerCAmelCase__        =  self.sp_model.serialized_model_proto()
                        fi.write(snake_case__				)

            return (str(snake_case__				), str(snake_case__				))







def 		_UpperCamelCase				(      UpperCamelCase_    :   Optional[Any]      ,    UpperCamelCase_    :   Optional[Any]						)					->				sentencepiece.SentencePieceProcessor:

      """simple docstring"""



      lowerCAmelCase__        =  sentencepiece.SentencePieceProcessor(**__lowerCAmelCase						)
      spm.Load(str(__lowerCAmelCase						)						)
      return spm







def 		_UpperCamelCase				(      UpperCamelCase_    :   List[Any]						)					->				Union[Dict, List]:

      """simple docstring"""



      with open(__lowerCAmelCase      ,    'r'						) as f:
            return json.load(__lowerCAmelCase						)







def 		_UpperCamelCase				(      UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   Tuple						)					->				None:

      """simple docstring"""



      with open(__lowerCAmelCase      ,    'w'						) as f:
            json.dump(__lowerCAmelCase      ,    __lowerCAmelCase      ,    indent=2						)













import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


__snake_case      :      Any             =       logging.get_logger(__name__)


class       __SCREAMING_SNAKE_CASE   (				__lowercase):
						_SCREAMING_SNAKE_CASE		:    Tuple						=   ['''input_features''', '''is_longer''']






						def __init__(		self       ,  _UpperCamelCase=64       ,  _UpperCamelCase=4_80_00       ,  _UpperCamelCase=4_80       ,  _UpperCamelCase=10       ,  _UpperCamelCase=10_24       ,  _UpperCamelCase=0.0       ,  _UpperCamelCase=False       ,  _UpperCamelCase = 0       ,  _UpperCamelCase = 1_40_00       ,  _UpperCamelCase = None       ,  _UpperCamelCase = "fusion"       ,  _UpperCamelCase = "repeatpad"       ,  **_UpperCamelCase       ,  ):






												"""simple docstring"""
												super().__init__(
												    feature_size=_UpperCamelCase       ,  sampling_rate=_UpperCamelCase       ,  padding_value=_UpperCamelCase       ,  return_attention_mask=_UpperCamelCase       ,  **_UpperCamelCase       ,  )
												lowerCAmelCase__        =  top_db
												lowerCAmelCase__        =  truncation
												lowerCAmelCase__        =  padding
												lowerCAmelCase__        =  fft_window_size
												lowerCAmelCase__        =  (fft_window_size >> 1) + 1
												lowerCAmelCase__        =  hop_length
												lowerCAmelCase__        =  max_length_s
												lowerCAmelCase__        =  max_length_s * sampling_rate
												lowerCAmelCase__        =  sampling_rate
												lowerCAmelCase__        =  frequency_min
												lowerCAmelCase__        =  frequency_max
												lowerCAmelCase__        =  mel_filter_bank(
												    num_frequency_bins=self.nb_frequency_bins       ,  num_mel_filters=_UpperCamelCase       ,  min_frequency=_UpperCamelCase       ,  max_frequency=_UpperCamelCase       ,  sampling_rate=_UpperCamelCase       ,  norm=_UpperCamelCase       ,  mel_scale='htk'       ,  )
												lowerCAmelCase__        =  mel_filter_bank(
												    num_frequency_bins=self.nb_frequency_bins       ,  num_mel_filters=_UpperCamelCase       ,  min_frequency=_UpperCamelCase       ,  max_frequency=_UpperCamelCase       ,  sampling_rate=_UpperCamelCase       ,  norm='slaney'       ,  mel_scale='slaney'       ,  )






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  copy.deepcopy(self.__dict__				)
												lowerCAmelCase__        =  self.__class__.__name__
												if "mel_filters" in output:
																		del output["mel_filters"]
												if "mel_filters_slaney" in output:
																		del output["mel_filters_slaney"]
												return output






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase = None				):






												"""simple docstring"""
												lowerCAmelCase__        =  spectrogram(
												    _UpperCamelCase       ,  window_function(self.fft_window_size       ,  'hann'				)       ,  frame_length=self.fft_window_size       ,  hop_length=self.hop_length       ,  power=2.0       ,  mel_filters=_UpperCamelCase       ,  log_mel='dB'       ,  )
												return log_mel_spectrogram.T






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												lowerCAmelCase__        =  np.array_split(list(range(0       ,  total_frames - chunk_frames + 1				)				)       ,  3				)
												if len(ranges[1]				) == 0:
																		# if the audio is too short, we just use the first chunk
																		lowerCAmelCase__        =  [0]
												if len(ranges[2]				) == 0:
																		# if the audio is too short, we just use the first chunk
																		lowerCAmelCase__        =  [0]
												# randomly choose index for each part
												lowerCAmelCase__        =  np.random.choice(ranges[0]				)
												lowerCAmelCase__        =  np.random.choice(ranges[1]				)
												lowerCAmelCase__        =  np.random.choice(ranges[2]				)

												lowerCAmelCase__        =  mel[idx_front : idx_front + chunk_frames, :]
												lowerCAmelCase__        =  mel[idx_middle : idx_middle + chunk_frames, :]
												lowerCAmelCase__        =  mel[idx_back : idx_back + chunk_frames, :]

												lowerCAmelCase__        =  torch.tensor(mel[None, None, :]				)
												lowerCAmelCase__        =  torch.nn.functional.interpolate(
												    _UpperCamelCase       ,  size=[chunk_frames, 64]       ,  mode='bilinear'       ,  align_corners=_UpperCamelCase				)
												lowerCAmelCase__        =  mel_shrink[0][0].numpy()
												lowerCAmelCase__        =  np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]       ,  axis=0				)
												return mel_fusion






						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				):






												"""simple docstring"""
												if waveform.shape[0] > max_length:
																		if truncation == "rand_trunc":
																								lowerCAmelCase__        =  True
																								# random crop to max_length (for compatibility) -> this should be handled by self.pad
																								lowerCAmelCase__        =  len(_UpperCamelCase				) - max_length
																								lowerCAmelCase__        =  np.random.randint(0       ,  overflow + 1				)
																								lowerCAmelCase__        =  waveform[idx : idx + max_length]
																								lowerCAmelCase__        =  self._np_extract_fbank_features(_UpperCamelCase       ,  self.mel_filters_slaney				)[None, :]
																		elif truncation == "fusion":
																								lowerCAmelCase__        =  self._np_extract_fbank_features(_UpperCamelCase       ,  self.mel_filters				)
																								lowerCAmelCase__        =  max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
																								lowerCAmelCase__        =  mel.shape[0]
																								if chunk_frames == total_frames:
																														# there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
																														# In this case, we just use the whole audio.
																														lowerCAmelCase__        =  np.stack([mel, mel, mel, mel]       ,  axis=0				)
																														lowerCAmelCase__        =  False
																								else:
																														lowerCAmelCase__        =  self._random_mel_fusion(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)
																														lowerCAmelCase__        =  True
																		else:
																								raise NotImplementedError(F"data_truncating {truncation} not implemented"				)

												else:
																		lowerCAmelCase__        =  False
																		# only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
																		if waveform.shape[0] < max_length:
																								if padding == "repeat":
																														lowerCAmelCase__        =  int(max_length / len(_UpperCamelCase				)				)
																														lowerCAmelCase__        =  np.stack(np.tile(_UpperCamelCase       ,  n_repeat + 1				)				)[:max_length]
																								if padding == "repeatpad":
																														lowerCAmelCase__        =  int(max_length / len(_UpperCamelCase				)				)
																														lowerCAmelCase__        =  np.stack(np.tile(_UpperCamelCase       ,  _UpperCamelCase				)				)
																								lowerCAmelCase__        =  np.pad(_UpperCamelCase       ,  (0, max_length - waveform.shape[0])       ,  mode='constant'       ,  constant_values=0				)

																		if truncation == "fusion":
																								lowerCAmelCase__        =  self._np_extract_fbank_features(_UpperCamelCase       ,  self.mel_filters				)
																								lowerCAmelCase__        =  np.stack([input_mel, input_mel, input_mel, input_mel]       ,  axis=0				)
																		else:
																								lowerCAmelCase__        =  self._np_extract_fbank_features(_UpperCamelCase       ,  self.mel_filters_slaney				)[None, :]

												return input_mel, longer







						def __call__(		self       ,  _UpperCamelCase       ,  _UpperCamelCase = None       ,  _UpperCamelCase = None       ,  _UpperCamelCase = None       ,  _UpperCamelCase = None       ,  _UpperCamelCase = None       ,  **_UpperCamelCase       ,  ):






												"""simple docstring"""
												lowerCAmelCase__        =  truncation if truncation is not None else self.truncation
												lowerCAmelCase__        =  padding if padding else self.padding

												if sampling_rate is not None:
																		if sampling_rate != self.sampling_rate:
																								raise ValueError(
																								    F"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
																								    F" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
																								    F" was sampled with {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.'				)

												lowerCAmelCase__        =  isinstance(_UpperCamelCase       ,  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}"				)
												lowerCAmelCase__        =  is_batched_numpy or (
												    isinstance(_UpperCamelCase       ,  (list, tuple)				) and (isinstance(raw_speech[0]       ,  (np.ndarray, tuple, list)				))
												)

												if is_batched:
																		lowerCAmelCase__        =  [np.asarray(_UpperCamelCase       ,  dtype=np.floataa				) for speech in raw_speech]
												elif not is_batched and not isinstance(_UpperCamelCase       ,  np.ndarray				):
																		lowerCAmelCase__        =  np.asarray(_UpperCamelCase       ,  dtype=np.floataa				)
												elif isinstance(_UpperCamelCase       ,  np.ndarray				) and raw_speech.dtype is np.dtype(np.floataa				):
																		lowerCAmelCase__        =  raw_speech.astype(np.floataa				)

												# always return batch
												if not is_batched:
																		lowerCAmelCase__        =  [np.asarray(_UpperCamelCase				)]

												# convert to mel spectrogram, truncate and pad if needed.
												lowerCAmelCase__        =  [
												    self._get_input_mel(_UpperCamelCase       ,  max_length if max_length else self.nb_max_samples       ,  _UpperCamelCase       ,  _UpperCamelCase				)
												    for waveform in raw_speech
												]

												lowerCAmelCase__        =  []
												lowerCAmelCase__        =  []
												for mel, longer in padded_inputs:
																		input_mel.append(_UpperCamelCase				)
																		is_longer.append(_UpperCamelCase				)

												if truncation == "fusion" and sum(_UpperCamelCase				) == 0:
																		# if no audio is longer than 10s, then randomly select one audio to be longer
																		lowerCAmelCase__        =  np.random.randint(0       ,  len(_UpperCamelCase				)				)
																		lowerCAmelCase__        =  True

												if isinstance(input_mel[0]       ,  _UpperCamelCase				):
																		lowerCAmelCase__        =  [np.asarray(_UpperCamelCase       ,  dtype=np.floataa				) for feature in input_mel]

												# is_longer is a list of bool
												lowerCAmelCase__        =  [[longer] for longer in is_longer]

												lowerCAmelCase__        =  {'input_features': input_mel, 'is_longer': is_longer}
												lowerCAmelCase__        =  BatchFeature(_UpperCamelCase				)

												if return_tensors is not None:
																		lowerCAmelCase__        =  input_features.convert_to_tensors(_UpperCamelCase				)

												return input_features





