


"""simple docstring"""


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

A  :     Optional[Any]										=    random.Random()


def _lowerCamelCase			(	_UpperCamelCase					,				_UpperCamelCase=1.0					,				_UpperCamelCase=None					,				_UpperCamelCase=None   ):

			'''simple docstring'''




			if rng is None:
						__lowerCAmelCase       					=    global_rng

			__lowerCAmelCase       					=    []
			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 						_UpperCamelCase						(     unittest.TestCase	):
							'''simple docstring'''







							def __init__(  self					,     __a					,     __a=7					,     __a=4_00					,     __a=20_00					,     __a=10					,     __a=1_60					,     __a=8					,     __a=0.0					,     __a=40_00					,     __a=False					,     __a=True					,     ):
										__lowerCAmelCase       					=    parent
										__lowerCAmelCase       					=    batch_size
										__lowerCAmelCase       					=    min_seq_length
										__lowerCAmelCase       					=    max_seq_length
										__lowerCAmelCase       					=    (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
										__lowerCAmelCase       					=    padding_value
										__lowerCAmelCase       					=    sampling_rate
										__lowerCAmelCase       					=    return_attention_mask
										__lowerCAmelCase       					=    do_normalize
										__lowerCAmelCase       					=    feature_size
										__lowerCAmelCase       					=    chunk_length
										__lowerCAmelCase       					=    hop_length

							def     snake_case       (  self			):
										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     snake_case       (  self					,     __a=False					,     __a=False			):
										def _flatten(__a			):
													return list(itertools.chain(*__a			)			)

										if equal_length:
													__lowerCAmelCase       					=    [floats_list((self.max_seq_length, self.feature_size)			) for _ in range(self.batch_size			)]
										else:
													# make sure that inputs increase in size
													__lowerCAmelCase       					=    [
													    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       					=    [np.asarray(__a			) for x in speech_inputs]
										return speech_inputs
@require_torch
@require_torchaudio
class 						_UpperCamelCase						(     lowerCAmelCase__ ,unittest.TestCase	):
							'''simple docstring'''







							__UpperCAmelCase		:					Optional[int]          =WhisperFeatureExtractor if is_speech_available() else None

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    WhisperFeatureExtractionTester(self			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.feature_extraction_class(**self.feat_extract_dict			)

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

										__lowerCAmelCase       					=    feat_extract_first.to_dict()
										__lowerCAmelCase       					=    feat_extract_second.to_dict()
										__lowerCAmelCase       					=    feat_extract_first.mel_filters
										__lowerCAmelCase       					=    feat_extract_second.mel_filters
										self.assertTrue(np.allclose(__a					,     __a			)			)
										self.assertEqual(__a					,     __a			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.feature_extraction_class(**self.feat_extract_dict			)

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

										__lowerCAmelCase       					=    feat_extract_first.to_dict()
										__lowerCAmelCase       					=    feat_extract_second.to_dict()
										__lowerCAmelCase       					=    feat_extract_first.mel_filters
										__lowerCAmelCase       					=    feat_extract_second.mel_filters
										self.assertTrue(np.allclose(__a					,     __a			)			)
										self.assertEqual(__a					,     __a			)

							def     snake_case       (  self			):
										# Tests that all call wrap to encode_plus and batch_encode_plus
										__lowerCAmelCase       					=    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()			)
										# create three inputs of length 800, 1000, and 1200
										__lowerCAmelCase       					=    [floats_list((1, x)			)[0] for x in range(8_00					,     14_00					,     2_00			)]
										__lowerCAmelCase       					=    [np.asarray(__a			) for speech_input in speech_inputs]

										# Test feature size
										__lowerCAmelCase       					=    feature_extractor(__a					,     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
										__lowerCAmelCase       					=    feature_extractor(speech_inputs[0]					,     return_tensors="np"			).input_features
										__lowerCAmelCase       					=    feature_extractor(np_speech_inputs[0]					,     return_tensors="np"			).input_features
										self.assertTrue(np.allclose(__a					,     __a					,     atol=1e-3			)			)

										# Test batched
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										for enc_seq_a, enc_seq_a in zip(__a					,     __a			):
													self.assertTrue(np.allclose(__a					,     __a					,     atol=1e-3			)			)

										# Test 2-D numpy arrays are batched.
										__lowerCAmelCase       					=    [floats_list((1, x)			)[0] for x in (8_00, 8_00, 8_00)]
										__lowerCAmelCase       					=    np.asarray(__a			)
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										for enc_seq_a, enc_seq_a in zip(__a					,     __a			):
													self.assertTrue(np.allclose(__a					,     __a					,     atol=1e-3			)			)

										# Test truncation required
										__lowerCAmelCase       					=    [floats_list((1, x)			)[0] for x in range(2_00					,     (feature_extractor.n_samples + 5_00)					,     2_00			)]
										__lowerCAmelCase       					=    [np.asarray(__a			) for speech_input in speech_inputs]

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

										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="np"			).input_features
										for enc_seq_a, enc_seq_a in zip(__a					,     __a			):
													self.assertTrue(np.allclose(__a					,     __a					,     atol=1e-3			)			)

							def     snake_case       (  self			):
										import torch

										__lowerCAmelCase       					=    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()			)
										__lowerCAmelCase       					=    np.random.rand(1_00					,     32			).astype(np.floataa			)
										__lowerCAmelCase       					=    np_speech_inputs.tolist()

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

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

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

							def     snake_case       (  self			):
										# fmt: off
										__lowerCAmelCase       					=    torch.tensor(
										    [
										        0.1_1_9_3, -0.0_9_4_6, -0.1_0_9_8, -0.0_1_9_6, 0.0_2_2_5, -0.0_6_9_0, -0.1_7_3_6, 0.0_9_5_1,
										        0.0_9_7_1, -0.0_8_1_7, -0.0_7_0_2, 0.0_1_6_2, 0.0_2_6_0, 0.0_0_1_7, -0.0_1_9_2, -0.1_6_7_8,
										        0.0_7_0_9, -0.1_8_6_7, -0.0_6_5_5, -0.0_2_7_4, -0.0_2_3_4, -0.1_8_8_4, -0.0_5_1_6, -0.0_5_5_4,
										        -0.0_2_7_4, -0.1_4_2_5, -0.1_4_2_3, 0.0_8_3_7, 0.0_3_7_7, -0.0_8_5_4
										    ]			)
										# fmt: on

										__lowerCAmelCase       					=    self._load_datasamples(1			)
										__lowerCAmelCase       					=    WhisperFeatureExtractor()
										__lowerCAmelCase       					=    feature_extractor(__a					,     return_tensors="pt"			).input_features
										self.assertEqual(input_features.shape					,     (1, 80, 30_00)			)
										self.assertTrue(torch.allclose(input_features[0, 0, :30]					,     __a					,     atol=1e-4			)			)



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()			)
										__lowerCAmelCase       					=    self._load_datasamples(1			)[0]
										__lowerCAmelCase       					=    ((audio - audio.min()) / (audio.max() - audio.min())) * 6_55_35  # Rescale to [0, 65535] to show issue
										__lowerCAmelCase       					=    feat_extract.zero_mean_unit_var_norm([audio]					,     attention_mask=__a			)[0]

										self.assertTrue(np.all(np.mean(__a			) < 1e-3			)			)
										self.assertTrue(np.all(np.abs(np.var(__a			) - 1			) < 1e-3			)			)








"""simple docstring"""


import tensorflow as tf

from ...tf_utils import shape_list
class 						_UpperCamelCase						(     tf.keras.layers.Layer	):
							'''simple docstring'''







							def __init__(  self					,     __a					,     __a					,     __a					,     __a					,     __a=1					,     __a=False					,     **__a			):
										super().__init__(**__a			)

										__lowerCAmelCase       					=    vocab_size
										__lowerCAmelCase       					=    d_embed
										__lowerCAmelCase       					=    d_proj

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

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

										__lowerCAmelCase       					=    []
										__lowerCAmelCase       					=    []

							def     snake_case       (  self					,     __a			):
										if self.n_clusters > 0:
													__lowerCAmelCase       					=    self.add_weight(
													    shape=(self.n_clusters, self.d_embed)					,     initializer="zeros"					,     trainable=__a					,     name="cluster_weight"			)
													__lowerCAmelCase       					=    self.add_weight(
													    shape=(self.n_clusters,)					,     initializer="zeros"					,     trainable=__a					,     name="cluster_bias"			)

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

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

							@staticmethod
							def     snake_case       (  __a					,     __a					,     __a					,     __a=None			):
										__lowerCAmelCase       					=    x
										if proj is not None:
													__lowerCAmelCase       					=    tf.einsum("ibd,ed->ibe"					,     __a					,     __a			)
										return tf.einsum("ibd,nd->ibn"					,     __a					,     __a			) + b

							@staticmethod
							def     snake_case       (  __a					,     __a			):
										__lowerCAmelCase       					=    shape_list(__a			)
										__lowerCAmelCase       					=    tf.range(lp_size[0]					,     dtype=target.dtype			)
										__lowerCAmelCase       					=    tf.stack([r, target]					,     1			)
										return tf.gather_nd(__a					,     __a			)



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

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

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

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

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

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

										return out




