
'''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

snake_case_ :  str     =					random.Random()



def  __snake_case     (    _UpperCAmelCase					:       List[str],      _UpperCAmelCase					:       List[Any]=1.0,      _UpperCAmelCase					:       int=None,      _UpperCAmelCase					:       List[str]=None):
   if rng is None:
      UpperCamelCase       =  global_rng

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

   return values


class 						lowercase__	(  unittest.TestCase   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__=7 ,   lowerCamelCase__=4_0_0 ,   lowerCamelCase__=2_0_0_0 ,   lowerCamelCase__=2_0_4_8 ,   lowerCamelCase__=1_2_8 ,   lowerCamelCase__=1 ,   lowerCamelCase__=5_1_2 ,   lowerCamelCase__=3_0 ,   lowerCamelCase__=4_4_1_0_0 ,   ):

      '''simple docstring'''







      UpperCamelCase       =  parent
      UpperCamelCase       =  batch_size
      UpperCamelCase       =  min_seq_length
      UpperCamelCase       =  max_seq_length
      UpperCamelCase       =  (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
      UpperCamelCase       =  spectrogram_length
      UpperCamelCase       =  feature_size
      UpperCamelCase       =  num_audio_channels
      UpperCamelCase       =  hop_length
      UpperCamelCase       =  chunk_length
      UpperCamelCase       =  sampling_rate







   def       UpperCAmelCase				(			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       UpperCAmelCase				(			self ,   lowerCamelCase__=False ,   lowerCamelCase__=False   ):

      '''simple docstring'''







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

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




@require_torch
@require_torchaudio
class 						lowercase__	(  snake_case_,					unittest.TestCase   ):






   '''simple docstring'''







   _snake_case		   =							TvltFeatureExtractor







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  TvltFeatureExtractionTester(self   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.feature_extraction_class(**self.feat_extract_dict   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''spectrogram_length'''   )   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''feature_size'''   )   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''num_audio_channels'''   )   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''hop_length'''   )   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''chunk_length'''   )   )
      self.assertTrue(hasattr(lowerCamelCase__ ,   '''sampling_rate'''   )   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.feature_extraction_class(**self.feat_extract_dict   )

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

      UpperCamelCase       =  feat_extract_first.to_dict()
      UpperCamelCase       =  feat_extract_second.to_dict()
      UpperCamelCase       =  dict_first.pop('''mel_filters'''   )
      UpperCamelCase       =  dict_second.pop('''mel_filters'''   )
      self.assertTrue(np.allclose(lowerCamelCase__ ,   lowerCamelCase__   )   )
      self.assertEqual(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.feature_extraction_class(**self.feat_extract_dict   )

      with tempfile.TemporaryDirectory() as tmpdirname:
         UpperCamelCase       =  os.path.join(lowerCamelCase__ ,   '''feat_extract.json'''   )
         feat_extract_first.to_json_file(lowerCamelCase__   )
         UpperCamelCase       =  self.feature_extraction_class.from_json_file(lowerCamelCase__   )

      UpperCamelCase       =  feat_extract_first.to_dict()
      UpperCamelCase       =  feat_extract_second.to_dict()
      UpperCamelCase       =  dict_first.pop('''mel_filters'''   )
      UpperCamelCase       =  dict_second.pop('''mel_filters'''   )
      self.assertTrue(np.allclose(lowerCamelCase__ ,   lowerCamelCase__   )   )
      self.assertEqual(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self.feature_extraction_class(**self.feat_extract_dict   )

      # create three inputs of length 800, 1000, and 1200
      UpperCamelCase       =  [floats_list((1, x)   )[0] for x in range(8_0_0 ,   1_4_0_0 ,   2_0_0   )]
      UpperCamelCase       =  [np.asarray(lowerCamelCase__   ) for speech_input in speech_inputs]

      # Test not batched input
      UpperCamelCase       =  feature_extractor(np_speech_inputs[0] ,   return_tensors='''np''' ,   sampling_rate=4_4_1_0_0   ).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
      UpperCamelCase       =  feature_extractor(lowerCamelCase__ ,   return_tensors='''np''' ,   sampling_rate=4_4_1_0_0   ).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
      UpperCamelCase       =  feature_extractor(
          lowerCamelCase__ ,   return_tensors='''np''' ,   sampling_rate=4_4_1_0_0 ,   mask_audio=lowerCamelCase__   ).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.
      UpperCamelCase       =  [floats_list((1, x)   )[0] for x in (8_0_0, 8_0_0, 8_0_0)]
      UpperCamelCase       =  np.asarray(lowerCamelCase__   )
      UpperCamelCase       =  feature_extractor(lowerCamelCase__ ,   return_tensors='''np''' ,   sampling_rate=4_4_1_0_0   ).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       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







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

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







   def       UpperCAmelCase				(			self   ):

      '''simple docstring'''







      UpperCamelCase       =  self._load_datasamples(1   )
      UpperCamelCase       =  TvltFeatureExtractor()
      UpperCamelCase       =  feature_extractor(lowerCamelCase__ ,   return_tensors='''pt'''   ).audio_values

      self.assertEquals(audio_values.shape ,   (1, 1, 1_9_2, 1_2_8)   )

      UpperCamelCase       =  torch.tensor([[-0.3032, -0.2708], [-0.4434, -0.4007]]   )
      self.assertTrue(torch.allclose(audio_values[0, 0, :2, :2] ,   lowerCamelCase__ ,   atol=1e-4   )   )







'''simple docstring'''


import tensorflow as tf

from ...tf_utils import shape_list


class 						lowercase__	(  tf.keras.layers.Layer   ):






   '''simple docstring'''







   def __init__(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=1 ,   lowerCamelCase__=False ,   **lowerCamelCase__   ):

      '''simple docstring'''







      super().__init__(**lowerCamelCase__   )

      UpperCamelCase       =  vocab_size
      UpperCamelCase       =  d_embed
      UpperCamelCase       =  d_proj

      UpperCamelCase       =  cutoffs + [vocab_size]
      UpperCamelCase       =  [0] + self.cutoffs
      UpperCamelCase       =  div_val

      UpperCamelCase       =  self.cutoffs[0]
      UpperCamelCase       =  len(self.cutoffs   ) - 1
      UpperCamelCase       =  self.shortlist_size + self.n_clusters
      UpperCamelCase       =  keep_order

      UpperCamelCase       =  []
      UpperCamelCase       =  []







   def       UpperCAmelCase				(			self ,   lowerCamelCase__   ):

      '''simple docstring'''







      if self.n_clusters > 0:
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_weight'''   )
         UpperCamelCase       =  self.add_weight(
             shape=(self.n_clusters,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name='''cluster_bias'''   )

      if self.div_val == 1:
         for i in range(len(self.cutoffs   )   ):
            if self.d_proj != self.d_embed:
               UpperCamelCase       =  self.add_weight(
                   shape=(self.d_embed, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}' ,   )
               self.out_projs.append(lowerCamelCase__   )
            else:
               self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size, self.d_embed) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(self.vocab_size,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      else:
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            UpperCamelCase       =  self.d_embed // (self.div_val**i)

            UpperCamelCase       =  self.add_weight(
                shape=(d_emb_i, self.d_proj) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_projs_._{i}'   )
            self.out_projs.append(lowerCamelCase__   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx, d_emb_i) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._weight' ,   )
            UpperCamelCase       =  self.add_weight(
                shape=(r_idx - l_idx,) ,   initializer='''zeros''' ,   trainable=lowerCamelCase__ ,   name=f'out_layers_._{i}_._bias' ,   )
            self.out_layers.append((weight, bias)   )
      super().build(lowerCamelCase__   )







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=None   ):

      '''simple docstring'''







      UpperCamelCase       =  x
      if proj is not None:
         UpperCamelCase       =  tf.einsum('''ibd,ed->ibe''' ,   lowerCamelCase__ ,   lowerCamelCase__   )
      return tf.einsum('''ibd,nd->ibn''' ,   lowerCamelCase__ ,   lowerCamelCase__   ) + b







   @staticmethod
   def       UpperCAmelCase				(			lowerCamelCase__ ,   lowerCamelCase__   ):

      '''simple docstring'''







      UpperCamelCase       =  shape_list(lowerCamelCase__   )
      UpperCamelCase       =  tf.range(lp_size[0] ,   dtype=target.dtype   )
      UpperCamelCase       =  tf.stack([r, target] ,   1   )
      return tf.gather_nd(lowerCamelCase__ ,   lowerCamelCase__   )







   def       UpperCAmelCase				(			self ,   lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__=True ,   lowerCamelCase__=False   ):

      '''simple docstring'''







      UpperCamelCase       =  0
      if self.n_clusters == 0:
         UpperCamelCase       =  self._logit(lowerCamelCase__ ,   self.out_layers[0][0] ,   self.out_layers[0][1] ,   self.out_projs[0]   )
         if target is not None:
            UpperCamelCase       =  tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__ ,   logits=lowerCamelCase__   )
         UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__ ,   axis=-1   )
      else:
         UpperCamelCase       =  shape_list(lowerCamelCase__   )
         UpperCamelCase       =  []
         UpperCamelCase       =  tf.zeros(hidden_sizes[:2]   )
         for i in range(len(self.cutoffs   )   ):
            UpperCamelCase     , UpperCamelCase       =  self.cutoff_ends[i], self.cutoff_ends[i + 1]
            if target is not None:
               UpperCamelCase       =  (target >= l_idx) & (target < r_idx)
               UpperCamelCase       =  tf.where(lowerCamelCase__   )
               UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   ) - l_idx

            if self.div_val == 1:
               UpperCamelCase       =  self.out_layers[0][0][l_idx:r_idx]
               UpperCamelCase       =  self.out_layers[0][1][l_idx:r_idx]
            else:
               UpperCamelCase       =  self.out_layers[i][0]
               UpperCamelCase       =  self.out_layers[i][1]

            if i == 0:
               UpperCamelCase       =  tf.concat([cur_W, self.cluster_weight] ,   0   )
               UpperCamelCase       =  tf.concat([cur_b, self.cluster_bias] ,   0   )

               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[0]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               out.append(head_logprob[..., : self.cutoffs[0]]   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
            else:
               UpperCamelCase       =  self._logit(lowerCamelCase__ ,   lowerCamelCase__ ,   lowerCamelCase__ ,   self.out_projs[i]   )
               UpperCamelCase       =  tf.nn.log_softmax(lowerCamelCase__   )
               UpperCamelCase       =  self.cutoffs[0] + i - 1  # No probability for the head cluster
               UpperCamelCase       =  head_logprob[..., cluster_prob_idx, None] + tail_logprob
               out.append(lowerCamelCase__   )
               if target is not None:
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  tf.boolean_mask(lowerCamelCase__ ,   lowerCamelCase__   )
                  UpperCamelCase       =  self._gather_logprob(lowerCamelCase__ ,   lowerCamelCase__   )
                  cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
            if target is not None:
               loss += tf.scatter_nd(lowerCamelCase__ ,   -cur_logprob ,   shape_list(lowerCamelCase__   )   )
         UpperCamelCase       =  tf.concat(lowerCamelCase__ ,   axis=-1   )

      if target is not None:
         if return_mean:
            UpperCamelCase       =  tf.reduce_mean(lowerCamelCase__   )
         # Add the training-time loss value to the layer using `self.add_loss()`.
         self.add_loss(lowerCamelCase__   )

         # Log the loss as a metric (we could log arbitrary metrics,
         # including different metrics for training and inference.
         self.add_metric(lowerCamelCase__ ,   name=self.name ,   aggregation='''mean''' if return_mean else ''''''   )

      return out





