





from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers.models.distilbert.modeling_tf_distilbert import (
         TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFDistilBertForMaskedLM,
         TFDistilBertForMultipleChoice,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertModel,
     )







class 				__lowerCAmelCase       :







      def __init__(    self		,      lowerCAmelCase__		,      )   ->						int:



        '''simple docstring'''




        a__	:					Optional[int]										=parent
        a__	:					Union[str, Any]										=1_3
        a__	:					Tuple										=7
        a__	:					List[str]										=True
        a__	:					Union[str, Any]										=True
        a__	:					Optional[Any]										=False
        a__	:					List[Any]										=True
        a__	:					Optional[int]										=9_9
        a__	:					Any										=3_2
        a__	:					List[str]										=2
        a__	:					List[Any]										=4
        a__	:					List[Any]										=3_7
        a__	:					Tuple										="gelu"
        a__	:					Optional[Any]										=0.1
        a__	:					Dict										=0.1
        a__	:					int										=5_1_2
        a__	:					str										=1_6
        a__	:					List[Any]										=2
        a__	:					int										=0.02
        a__	:					List[Any]										=3
        a__	:					Any										=4
        a__	:					Optional[int]										=None







      def    _lowercase  (    self							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=ids_tensor([self.batch_size, self.seq_length]		,      self.vocab_size							)

        a__	:					str										=None
        if self.use_input_mask:
          a__	:					Optional[Any]										=random_attention_mask([self.batch_size, self.seq_length]							)

        a__	:					List[str]										=None
        a__	:					str										=None
        a__	:					Dict										=None
        if self.use_labels:
          a__	:					Optional[int]										=ids_tensor([self.batch_size]		,      self.type_sequence_label_size							)
          a__	:					List[Any]										=ids_tensor([self.batch_size, self.seq_length]		,      self.num_labels							)
          a__	:					Any										=ids_tensor([self.batch_size]		,      self.num_choices							)

        a__	:					List[str]										=DistilBertConfig(
            vocab_size=self.vocab_size		,      dim=self.hidden_size		,      n_layers=self.num_hidden_layers		,      n_heads=self.num_attention_heads		,      hidden_dim=self.intermediate_size		,      hidden_act=self.hidden_act		,      dropout=self.hidden_dropout_prob		,      attention_dropout=self.attention_probs_dropout_prob		,      max_position_embeddings=self.max_position_embeddings		,      initializer_range=self.initializer_range		,      )

        return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					int										=TFDistilBertModel(config=lowerCAmelCase__							)
        a__	:					Tuple										={"input_ids": input_ids, "attention_mask": input_mask}

        a__	:					List[Any]										=model(lowerCAmelCase__							)

        a__	:					Tuple										=[input_ids, input_mask]

        a__	:					Tuple										=model(lowerCAmelCase__							)

        self.parent.assertEqual(result.last_hidden_state.shape		,      (self.batch_size, self.seq_length, self.hidden_size)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Union[str, Any]										=TFDistilBertForMaskedLM(config=lowerCAmelCase__							)
        a__	:					Dict										={"input_ids": input_ids, "attention_mask": input_mask}
        a__	:					List[str]										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.vocab_size)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						List[Any]:



        '''simple docstring'''




        a__	:					List[str]										=TFDistilBertForQuestionAnswering(config=lowerCAmelCase__							)
        a__	:					Union[str, Any]										={
            "input_ids": input_ids,
            "attention_mask": input_mask,
        }
        a__	:					int										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.start_logits.shape		,      (self.batch_size, self.seq_length)							)
        self.parent.assertEqual(result.end_logits.shape		,      (self.batch_size, self.seq_length)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Tuple										=self.num_labels
        a__	:					Optional[int]										=TFDistilBertForSequenceClassification(lowerCAmelCase__							)
        a__	:					List[Any]										={"input_ids": input_ids, "attention_mask": input_mask}
        a__	:					List[Any]										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_labels)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Union[str, Any]										=self.num_choices
        a__	:					Dict										=TFDistilBertForMultipleChoice(lowerCAmelCase__							)
        a__	:					Dict										=tf.tile(tf.expand_dims(lowerCAmelCase__		,      1							)		,      (1, self.num_choices, 1)							)
        a__	:					Optional[int]										=tf.tile(tf.expand_dims(lowerCAmelCase__		,      1							)		,      (1, self.num_choices, 1)							)
        a__	:					Optional[int]										={
            "input_ids": multiple_choice_inputs_ids,
            "attention_mask": multiple_choice_input_mask,
        }
        a__	:					int										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.num_choices)							)







      def    _lowercase  (    self		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__		,      lowerCAmelCase__							)   ->						Optional[Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=self.num_labels
        a__	:					Dict										=TFDistilBertForTokenClassification(lowerCAmelCase__							)
        a__	:					str										={"input_ids": input_ids, "attention_mask": input_mask}
        a__	:					Tuple										=model(lowerCAmelCase__							)
        self.parent.assertEqual(result.logits.shape		,      (self.batch_size, self.seq_length, self.num_labels)							)







      def    _lowercase  (    self							)   ->						Dict:



        '''simple docstring'''




        a__	:					Dict										=self.prepare_config_and_inputs()
        ((a__)					,							(a__)					,							(a__)					,							(a__)					,							(a__)					,							(a__))	:					Dict										=config_and_inputs
        a__	:					List[str]										={"input_ids": input_ids, "attention_mask": input_mask}
        return config, inputs_dict







@require_tf
class 				__lowerCAmelCase       (    UpperCamelCase__						,				UpperCamelCase__						,				unittest.TestCase):
      _lowercase				:					Any						=			(
          (
              TFDistilBertModel,
              TFDistilBertForMaskedLM,
              TFDistilBertForQuestionAnswering,
              TFDistilBertForSequenceClassification,
              TFDistilBertForTokenClassification,
              TFDistilBertForMultipleChoice,
          )
          if is_tf_available()
          else None
      )
      _lowercase				:					List[Any]						=			(
          {
              """feature-extraction""": TFDistilBertModel,
              """fill-mask""": TFDistilBertForMaskedLM,
              """question-answering""": TFDistilBertForQuestionAnswering,
              """text-classification""": TFDistilBertForSequenceClassification,
              """token-classification""": TFDistilBertForTokenClassification,
              """zero-shot""": TFDistilBertForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _lowercase				:					int						=			False
      _lowercase				:					List[str]						=			False







      def    _lowercase  (    self							)   ->						Tuple:



        '''simple docstring'''




        a__	:					Dict										=TFDistilBertModelTester(self							)
        a__	:					Optional[int]										=ConfigTester(self		,      config_class=lowerCAmelCase__		,      dim=3_7							)







      def    _lowercase  (    self							)   ->						int:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def    _lowercase  (    self							)   ->						Tuple:



        '''simple docstring'''




        a__	:					int										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Optional[int]										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					Dict										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						List[str]:



        '''simple docstring'''




        a__	:					List[str]										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						str:



        '''simple docstring'''




        a__	:					int										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase__							)







      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        a__	:					Union[str, Any]										=self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase__							)







      @slow
      def    _lowercase  (    self							)   ->						Union[str, Any]:



        '''simple docstring'''




        for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]							):
          a__	:					List[str]										=TFDistilBertModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)







@require_tf
class 				__lowerCAmelCase       (    unittest.TestCase):







      @slow
      def    _lowercase  (    self							)   ->						Optional[int]:



        '''simple docstring'''




        a__	:					Dict										=TFDistilBertModel.from_pretrained("distilbert-base-uncased"							)
        a__	:					Optional[int]										=tf.constant([[0, 1, 2, 3, 4, 5]]							)
        a__	:					Tuple										=model(lowerCAmelCase__							)[0]

        a__	:					List[str]										=[1, 6, 7_6_8]
        self.assertEqual(output.shape		,      lowerCAmelCase__							)

        a__	:					List[Any]										=tf.constant(
            [
                [
                    [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
                    [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
                    [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
                ]
            ]							)
        tf.debugging.assert_near(output[:, :3, :3]		,      lowerCAmelCase__		,      atol=1E-4							)











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


A__								=						logging.get_logger(__name__)


class       a							(			__lowerCamelCase				):
		__lowerCAmelCase	:    str											=					["""input_features""", """is_longer"""]



		def __init__(     self     :Union[str, Any]   ,__lowercase     :str=6_4   ,__lowercase     :Any=4_8_0_0_0   ,__lowercase     :List[Any]=4_8_0   ,__lowercase     :Optional[int]=1_0   ,__lowercase     :Optional[int]=1_0_2_4   ,__lowercase     :int=0.0   ,__lowercase     :List[Any]=False   ,__lowercase     :float = 0   ,__lowercase     :float = 1_4_0_0_0   ,__lowercase     :int = None   ,__lowercase     :str = "fusion"   ,__lowercase     :str = "repeatpad"   ,**__lowercase     :List[Any]   ,):
						super().__init__(
						    feature_size=__lowercase   ,sampling_rate=__lowercase   ,padding_value=__lowercase   ,return_attention_mask=__lowercase   ,**__lowercase   ,)
						snake_case__	:      Optional[Any]     		=	top_db
						snake_case__	:      Tuple     		=	truncation
						snake_case__	:      Tuple     		=	padding
						snake_case__	:      List[Any]     		=	fft_window_size
						snake_case__	:      List[Any]     		=	(fft_window_size >> 1) + 1
						snake_case__	:      str     		=	hop_length
						snake_case__	:      Dict     		=	max_length_s
						snake_case__	:      List[str]     		=	max_length_s * sampling_rate
						snake_case__	:      List[Any]     		=	sampling_rate
						snake_case__	:      Any     		=	frequency_min
						snake_case__	:      Dict     		=	frequency_max
						snake_case__	:      Tuple     		=	mel_filter_bank(
						    num_frequency_bins=self.nb_frequency_bins   ,num_mel_filters=__lowercase   ,min_frequency=__lowercase   ,max_frequency=__lowercase   ,sampling_rate=__lowercase   ,norm=__lowercase   ,mel_scale='''htk'''   ,)
						snake_case__	:      Dict     		=	mel_filter_bank(
						    num_frequency_bins=self.nb_frequency_bins   ,num_mel_filters=__lowercase   ,min_frequency=__lowercase   ,max_frequency=__lowercase   ,sampling_rate=__lowercase   ,norm='''slaney'''   ,mel_scale='''slaney'''   ,)



		def 	__lowerCamelCase			(     self     :int  ):
						snake_case__	:      Dict     		=	copy.deepcopy(self.__dict__  )
						snake_case__	:      Optional[int]     		=	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 	__lowerCamelCase			(     self     :Union[str, Any]   ,__lowercase     :np.array   ,__lowercase     :Optional[np.array] = None  ):
						snake_case__	:      List[Any]     		=	spectrogram(
						    __lowercase   ,window_function(self.fft_window_size   ,'''hann'''  )   ,frame_length=self.fft_window_size   ,hop_length=self.hop_length   ,power=2.0   ,mel_filters=__lowercase   ,log_mel='''dB'''   ,)
						return log_mel_spectrogram.T



		def 	__lowerCamelCase			(     self     :Optional[Any]   ,__lowercase     :List[str]   ,__lowercase     :Tuple   ,__lowercase     :List[str]  ):
						snake_case__	:      Dict     		=	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
										snake_case__	:      Tuple     		=	[0]
						if len(ranges[2]  ) == 0:
										# if the audio is too short, we just use the first chunk
										snake_case__	:      str     		=	[0]
						# randomly choose index for each part
						snake_case__	:      Dict     		=	np.random.choice(ranges[0]  )
						snake_case__	:      Any     		=	np.random.choice(ranges[1]  )
						snake_case__	:      Dict     		=	np.random.choice(ranges[2]  )

						snake_case__	:      Optional[int]     		=	mel[idx_front : idx_front + chunk_frames, :]
						snake_case__	:      Optional[Any]     		=	mel[idx_middle : idx_middle + chunk_frames, :]
						snake_case__	:      List[str]     		=	mel[idx_back : idx_back + chunk_frames, :]

						snake_case__	:      Optional[Any]     		=	torch.tensor(mel[None, None, :]  )
						snake_case__	:      Any     		=	torch.nn.functional.interpolate(
						    __lowercase   ,size=[chunk_frames, 6_4]   ,mode='''bilinear'''   ,align_corners=__lowercase  )
						snake_case__	:      List[Any]     		=	mel_shrink[0][0].numpy()
						snake_case__	:      Union[str, Any]     		=	np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back]   ,axis=0  )
						return mel_fusion



		def 	__lowerCamelCase			(     self     :Any   ,__lowercase     :np.array   ,__lowercase     :str   ,__lowercase     :int   ,__lowercase     :List[str]  ):
						if waveform.shape[0] > max_length:
										if truncation == "rand_trunc":
														snake_case__	:      Tuple     		=	True
														# random crop to max_length (for compatibility) -> this should be handled by self.pad
														snake_case__	:      List[str]     		=	len(__lowercase  ) - max_length
														snake_case__	:      Any     		=	np.random.randint(0   ,overflow + 1  )
														snake_case__	:      Tuple     		=	waveform[idx : idx + max_length]
														snake_case__	:      Tuple     		=	self._np_extract_fbank_features(__lowercase   ,self.mel_filters_slaney  )[None, :]
										elif truncation == "fusion":
														snake_case__	:      str     		=	self._np_extract_fbank_features(__lowercase   ,self.mel_filters  )
														snake_case__	:      Union[str, Any]     		=	max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
														snake_case__	:      Union[str, Any]     		=	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.
																		snake_case__	:      Tuple     		=	np.stack([mel, mel, mel, mel]   ,axis=0  )
																		snake_case__	:      List[Any]     		=	False
														else:
																		snake_case__	:      List[Any]     		=	self._random_mel_fusion(__lowercase   ,__lowercase   ,__lowercase  )
																		snake_case__	:      Dict     		=	True
										else:
														raise NotImplementedError(F"""data_truncating {truncation} not implemented"""  )

						else:
										snake_case__	:      Optional[int]     		=	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":
																		snake_case__	:      List[str]     		=	int(max_length / len(__lowercase  )  )
																		snake_case__	:      List[str]     		=	np.stack(np.tile(__lowercase   ,n_repeat + 1  )  )[:max_length]
														if padding == "repeatpad":
																		snake_case__	:      Union[str, Any]     		=	int(max_length / len(__lowercase  )  )
																		snake_case__	:      List[str]     		=	np.stack(np.tile(__lowercase   ,__lowercase  )  )
														snake_case__	:      int     		=	np.pad(__lowercase   ,(0, max_length - waveform.shape[0])   ,mode='''constant'''   ,constant_values=0  )

										if truncation == "fusion":
														snake_case__	:      Tuple     		=	self._np_extract_fbank_features(__lowercase   ,self.mel_filters  )
														snake_case__	:      Optional[int]     		=	np.stack([input_mel, input_mel, input_mel, input_mel]   ,axis=0  )
										else:
														snake_case__	:      List[Any]     		=	self._np_extract_fbank_features(__lowercase   ,self.mel_filters_slaney  )[None, :]

						return input_mel, longer






		def __call__(     self     :Dict   ,__lowercase     :Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]]   ,__lowercase     :str = None   ,__lowercase     :Optional[str] = None   ,__lowercase     :Optional[int] = None   ,__lowercase     :Optional[int] = None   ,__lowercase     :Optional[Union[str, TensorType]] = None   ,**__lowercase     :Optional[int]   ,):
						snake_case__	:      Optional[int]     		=	truncation if truncation is not None else self.truncation
						snake_case__	:      Optional[int]     		=	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.'''  )

						snake_case__	:      List[str]     		=	isinstance(__lowercase   ,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__	:      Optional[int]     		=	is_batched_numpy or (
						    isinstance(__lowercase   ,(list, tuple)  ) and (isinstance(raw_speech[0]   ,(np.ndarray, tuple, list)  ))
						)

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

						# always return batch
						if not is_batched:
										snake_case__	:      Dict     		=	[np.asarray(__lowercase  )]

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

						snake_case__	:      Any     		=	[]
						snake_case__	:      Tuple     		=	[]
						for mel, longer in padded_inputs:
										input_mel.append(__lowercase  )
										is_longer.append(__lowercase  )

						if truncation == "fusion" and sum(__lowercase  ) == 0:
										# if no audio is longer than 10s, then randomly select one audio to be longer
										snake_case__	:      Optional[int]     		=	np.random.randint(0   ,len(__lowercase  )  )
										snake_case__	:      List[str]     		=	True

						if isinstance(input_mel[0]   ,__lowercase  ):
										snake_case__	:      Optional[int]     		=	[np.asarray(__lowercase   ,dtype=np.floataa  ) for feature in input_mel]

						# is_longer is a list of bool
						snake_case__	:      Dict     		=	[[longer] for longer in is_longer]

						snake_case__	:      Dict     		=	{'''input_features''': input_mel, '''is_longer''': is_longer}
						snake_case__	:      str     		=	BatchFeature(__lowercase  )

						if return_tensors is not None:
										snake_case__	:      List[str]     		=	input_features.convert_to_tensors(__lowercase  )

						return input_features

