




import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      DistilBertForMaskedLM,
      DistilBertForMultipleChoice,
      DistilBertForQuestionAnswering,
      DistilBertForSequenceClassification,
      DistilBertForTokenClassification,
      DistilBertModel,
  )
class 			_UpperCamelCase      (				__snake_case		):
      def __init__(			self   :Any ,      lowerCamelCase   :Any ,      lowerCamelCase   :Optional[Any]=13 ,      lowerCamelCase   :Dict=7 ,      lowerCamelCase   :Tuple=True ,      lowerCamelCase   :List[Any]=True ,      lowerCamelCase   :List[str]=False ,      lowerCamelCase   :Union[str, Any]=True ,      lowerCamelCase   :Optional[int]=99 ,      lowerCamelCase   :Dict=32 ,      lowerCamelCase   :Any=5 ,      lowerCamelCase   :Optional[int]=4 ,      lowerCamelCase   :List[str]=37 ,      lowerCamelCase   :Optional[int]="gelu" ,      lowerCamelCase   :Tuple=0.1 ,      lowerCamelCase   :str=0.1 ,      lowerCamelCase   :Dict=512 ,      lowerCamelCase   :Tuple=16 ,      lowerCamelCase   :Dict=2 ,      lowerCamelCase   :Optional[Any]=0.02 ,      lowerCamelCase   :Tuple=3 ,      lowerCamelCase   :Tuple=4 ,      lowerCamelCase   :List[str]=None ,      )     ->			Tuple:
       UpperCAmelCase__							 =						parent
       UpperCAmelCase__							 =						batch_size
       UpperCAmelCase__							 =						seq_length
       UpperCAmelCase__							 =						is_training
       UpperCAmelCase__							 =						use_input_mask
       UpperCAmelCase__							 =						use_token_type_ids
       UpperCAmelCase__							 =						use_labels
       UpperCAmelCase__							 =						vocab_size
       UpperCAmelCase__							 =						hidden_size
       UpperCAmelCase__							 =						num_hidden_layers
       UpperCAmelCase__							 =						num_attention_heads
       UpperCAmelCase__							 =						intermediate_size
       UpperCAmelCase__							 =						hidden_act
       UpperCAmelCase__							 =						hidden_dropout_prob
       UpperCAmelCase__							 =						attention_probs_dropout_prob
       UpperCAmelCase__							 =						max_position_embeddings
       UpperCAmelCase__							 =						type_vocab_size
       UpperCAmelCase__							 =						type_sequence_label_size
       UpperCAmelCase__							 =						initializer_range
       UpperCAmelCase__							 =						num_labels
       UpperCAmelCase__							 =						num_choices
       UpperCAmelCase__							 =						scope
      def 	UpperCAmelCase_						(			self   :Union[str, Any]      )     ->			int:
       UpperCAmelCase__							 =						ids_tensor([self.batch_size, self.seq_length] ,      self.vocab_size      )

       UpperCAmelCase__							 =						None
       if self.use_input_mask:
        UpperCAmelCase__							 =						random_attention_mask([self.batch_size, self.seq_length]      )

       UpperCAmelCase__							 =						None
       UpperCAmelCase__							 =						None
       UpperCAmelCase__							 =						None
       if self.use_labels:
        UpperCAmelCase__							 =						ids_tensor([self.batch_size] ,      self.type_sequence_label_size      )
        UpperCAmelCase__							 =						ids_tensor([self.batch_size, self.seq_length] ,      self.num_labels      )
        UpperCAmelCase__							 =						ids_tensor([self.batch_size] ,      self.num_choices      )

       UpperCAmelCase__							 =						self.get_config()

       return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
      def 	UpperCAmelCase_						(			self   :str      )     ->			Tuple:
       return 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 ,      )
      def 	UpperCAmelCase_						(			self   :Dict ,      lowerCamelCase   :int ,      lowerCamelCase   :Tuple ,      lowerCamelCase   :List[str] ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :Dict ,      lowerCamelCase   :Any      )     ->			Tuple:
       UpperCAmelCase__							 =						DistilBertModel(config=_lowercase      )
       model.to(_lowercase      )
       model.eval()
       UpperCAmelCase__							 =						model(_lowercase ,      _lowercase      )
       UpperCAmelCase__							 =						model(_lowercase      )
       self.parent.assertEqual(result.last_hidden_state.shape ,      (self.batch_size, self.seq_length, self.hidden_size)      )
      def 	UpperCAmelCase_						(			self   :Dict ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :Union[str, Any] ,      lowerCamelCase   :str ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :Tuple ,      lowerCamelCase   :Optional[int]      )     ->			List[Any]:
       UpperCAmelCase__							 =						DistilBertForMaskedLM(config=_lowercase      )
       model.to(_lowercase      )
       model.eval()
       UpperCAmelCase__							 =						model(_lowercase ,      attention_mask=_lowercase ,      labels=_lowercase      )
       self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.vocab_size)      )
      def 	UpperCAmelCase_						(			self   :int ,      lowerCamelCase   :Any ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :str ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :int ,      lowerCamelCase   :Any      )     ->			List[str]:
       UpperCAmelCase__							 =						DistilBertForQuestionAnswering(config=_lowercase      )
       model.to(_lowercase      )
       model.eval()
       UpperCAmelCase__							 =						model(
           _lowercase ,      attention_mask=_lowercase ,      start_positions=_lowercase ,      end_positions=_lowercase      )
       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 	UpperCAmelCase_						(			self   :Dict ,      lowerCamelCase   :Tuple ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :Any ,      lowerCamelCase   :Optional[int] ,      lowerCamelCase   :Optional[Any] ,      lowerCamelCase   :Any      )     ->			Tuple:
       UpperCAmelCase__							 =						self.num_labels
       UpperCAmelCase__							 =						DistilBertForSequenceClassification(_lowercase      )
       model.to(_lowercase      )
       model.eval()
       UpperCAmelCase__							 =						model(_lowercase ,      attention_mask=_lowercase ,      labels=_lowercase      )
       self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_labels)      )
      def 	UpperCAmelCase_						(			self   :Dict ,      lowerCamelCase   :str ,      lowerCamelCase   :Optional[Any] ,      lowerCamelCase   :Union[str, Any] ,      lowerCamelCase   :Optional[Any] ,      lowerCamelCase   :str ,      lowerCamelCase   :Tuple      )     ->			Tuple:
       UpperCAmelCase__							 =						self.num_labels
       UpperCAmelCase__							 =						DistilBertForTokenClassification(config=_lowercase      )
       model.to(_lowercase      )
       model.eval()

       UpperCAmelCase__							 =						model(_lowercase ,      attention_mask=_lowercase ,      labels=_lowercase      )
       self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.seq_length, self.num_labels)      )
      def 	UpperCAmelCase_						(			self   :Optional[Any] ,      lowerCamelCase   :str ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :Optional[Any] ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :List[Any] ,      lowerCamelCase   :Optional[int]      )     ->			List[Any]:
       UpperCAmelCase__							 =						self.num_choices
       UpperCAmelCase__							 =						DistilBertForMultipleChoice(config=_lowercase      )
       model.to(_lowercase      )
       model.eval()
       UpperCAmelCase__							 =						input_ids.unsqueeze(1      ).expand(-1 ,      self.num_choices ,      -1      ).contiguous()
       UpperCAmelCase__							 =						input_mask.unsqueeze(1      ).expand(-1 ,      self.num_choices ,      -1      ).contiguous()
       UpperCAmelCase__							 =						model(
           _lowercase ,      attention_mask=_lowercase ,      labels=_lowercase ,      )
       self.parent.assertEqual(result.logits.shape ,      (self.batch_size, self.num_choices)      )







      def 	UpperCAmelCase_						(			self   :List[Any]      )     ->			str:
       UpperCAmelCase__							 =						self.prepare_config_and_inputs()
       (UpperCAmelCase__)							 =						config_and_inputs
       UpperCAmelCase__							 =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
       return config, inputs_dict



@require_torch
class 			_UpperCamelCase      (				__snake_case	,       __snake_case	,       unittest.TestCase		):
      UpperCAmelCase_        =		(
          (
              DistilBertModel,
              DistilBertForMaskedLM,
              DistilBertForMultipleChoice,
              DistilBertForQuestionAnswering,
              DistilBertForSequenceClassification,
              DistilBertForTokenClassification,
          )
          if is_torch_available()
          else None
      )
      UpperCAmelCase_        =		(
          {
              """feature-extraction""": DistilBertModel,
              """fill-mask""": DistilBertForMaskedLM,
              """question-answering""": DistilBertForQuestionAnswering,
              """text-classification""": DistilBertForSequenceClassification,
              """token-classification""": DistilBertForTokenClassification,
              """zero-shot""": DistilBertForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      UpperCAmelCase_        =		True
      UpperCAmelCase_        =		True
      UpperCAmelCase_        =		True
      UpperCAmelCase_        =		True
      def 	UpperCAmelCase_						(			self   :Optional[Any]      )     ->			Union[str, Any]:
       UpperCAmelCase__							 =						DistilBertModelTester(self      )
       UpperCAmelCase__							 =						ConfigTester(self ,      config_class=_lowercase ,      dim=37      )
      def 	UpperCAmelCase_						(			self   :str      )     ->			Optional[Any]:
       self.config_tester.run_common_tests()
      def 	UpperCAmelCase_						(			self   :List[str]      )     ->			Dict:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*_lowercase      )
      def 	UpperCAmelCase_						(			self   :Tuple      )     ->			Optional[Any]:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*_lowercase      )
      def 	UpperCAmelCase_						(			self   :Any      )     ->			List[str]:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*_lowercase      )
      def 	UpperCAmelCase_						(			self   :int      )     ->			str:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*_lowercase      )
      def 	UpperCAmelCase_						(			self   :int      )     ->			Optional[int]:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*_lowercase      )
      def 	UpperCAmelCase_						(			self   :Union[str, Any]      )     ->			int:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*_lowercase      )
      @slow
      def 	UpperCAmelCase_						(			self   :int      )     ->			int:

       for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							 =						DistilBertModel.from_pretrained(_lowercase      )
        self.assertIsNotNone(_lowercase      )







      @slow
      @require_torch_gpu
      def 	UpperCAmelCase_						(			self   :Tuple      )     ->			Any:
       UpperCAmelCase__							 =						self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
        # BertForMultipleChoice behaves incorrectly in JIT environments.
        if model_class == DistilBertForMultipleChoice:
         return

        UpperCAmelCase__							 =						True
        UpperCAmelCase__							 =						model_class(config=_lowercase      )

        UpperCAmelCase__							 =						self._prepare_for_class(_lowercase ,      _lowercase      )
        UpperCAmelCase__							 =						torch.jit.trace(
            _lowercase ,      (inputs_dict["input_ids"].to("cpu"      ), inputs_dict["attention_mask"].to("cpu"      ))      )

        with tempfile.TemporaryDirectory() as tmp:
         torch.jit.save(_lowercase ,      os.path.join(_lowercase ,      "traced_model.pt"      )      )
         UpperCAmelCase__							 =						torch.jit.load(os.path.join(_lowercase ,      "traced_model.pt"      ) ,      map_location=_lowercase      )
         loaded(inputs_dict["input_ids"].to(_lowercase      ) ,      inputs_dict["attention_mask"].to(_lowercase      )      )



@require_torch
class 			_UpperCamelCase      (				unittest.TestCase		):
      @slow
      def 	UpperCAmelCase_						(			self   :Dict      )     ->			List[str]:
       UpperCAmelCase__							 =						DistilBertModel.from_pretrained("distilbert-base-uncased"      )
       UpperCAmelCase__							 =						torch.tensor([[0, 345, 232, 328, 740, 140, 1695, 69, 6078, 1588, 2]]      )
       UpperCAmelCase__							 =						torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]      )
       with torch.no_grad():
        UpperCAmelCase__							 =						model(_lowercase ,      attention_mask=_lowercase      )[0]
       UpperCAmelCase__							 =						torch.Size((1, 11, 768)      )
       self.assertEqual(output.shape ,      _lowercase      )
       UpperCAmelCase__							 =						torch.tensor(
           [[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]]      )

       self.assertTrue(torch.allclose(output[:, 1:4, 1:4] ,      _lowercase ,      atol=1e-4      )      )











from typing import List, Optional, Union

import numpy as np

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


_lowerCAmelCase						:   Any		=   logging.get_logger(__name__)
class 			_UpperCamelCase      (				lowerCAmelCase		):
      UpperCAmelCase_        =		["""input_values""", """padding_mask"""]
      def __init__(			self   :Dict ,      lowerCamelCase   :int = 1 ,      lowerCamelCase   :int = 2_4000 ,      lowerCamelCase   :float = 0.0 ,      lowerCamelCase   :float = None ,      lowerCamelCase   :float = None ,      **lowerCamelCase   :Optional[Any] ,      )     ->			str:
       super().__init__(feature_size=lowerCamelCase ,      sampling_rate=lowerCamelCase ,      padding_value=lowerCamelCase ,      **lowerCamelCase      )
       UpperCAmelCase__							 =						chunk_length_s
       UpperCAmelCase__							 =						overlap
      @property
      def 	UpperCAmelCase_						(			self   :Optional[Any]      )     ->			Optional[int]:
       if self.chunk_length_s is None:
        return None
       else:
        return int(self.chunk_length_s * self.sampling_rate      )
      @property
      def 	UpperCAmelCase_						(			self   :str      )     ->			Optional[int]:
       if self.chunk_length_s is None or self.overlap is None:
        return None
       else:
        return max(1 ,      int((1.0 - self.overlap) * self.chunk_length      )      )







      def __call__(			self   :List[Any] ,      lowerCamelCase   :Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]] ,      lowerCamelCase   :Optional[Union[bool, str, PaddingStrategy]] = None ,      lowerCamelCase   :Optional[bool] = False ,      lowerCamelCase   :Optional[int] = None ,      lowerCamelCase   :Optional[Union[str, TensorType]] = None ,      lowerCamelCase   :Optional[int] = None ,      )     ->			BatchFeature:
       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 audio 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."      )

       if padding and truncation:
        raise ValueError("Both padding and truncation were set. Make sure you only set one."      )
       elif padding is None:
        # by default let's pad the inputs
        UpperCAmelCase__							 =						True

       UpperCAmelCase__							 =						bool(
           isinstance(lowerCamelCase ,      (list, tuple)      ) and (isinstance(raw_audio[0] ,      (np.ndarray, tuple, list)      ))      )

       if is_batched:
        UpperCAmelCase__							 =						[np.asarray(lowerCamelCase ,      dtype=np.floataa      ).T for audio in raw_audio]
       elif not is_batched and not isinstance(lowerCamelCase ,      np.ndarray      ):
        UpperCAmelCase__							 =						np.asarray(lowerCamelCase ,      dtype=np.floataa      )
       elif isinstance(lowerCamelCase ,      np.ndarray      ) and raw_audio.dtype is np.dtype(np.floataa      ):
        UpperCAmelCase__							 =						raw_audio.astype(np.floataa      )

       # always return batch
       if not is_batched:
        UpperCAmelCase__							 =						[np.asarray(lowerCamelCase      ).T]

       # verify inputs are valid
       for idx, example in enumerate(lowerCamelCase      ):
        if example.ndim > 2:
         raise ValueError(f'''Expected input shape (channels, length) but got shape {example.shape}'''      )
        if self.feature_size == 1 and example.ndim != 1:
         raise ValueError(f'''Expected mono audio but example has {example.shape[-1]} channels'''      )
        if self.feature_size == 2 and example.shape[-1] != 2:
         raise ValueError(f'''Expected stereo audio but example has {example.shape[-1]} channels'''      )

       UpperCAmelCase__							 =						None
       UpperCAmelCase__							 =						BatchFeature({"input_values": raw_audio}      )
       if self.chunk_stride is not None and self.chunk_length is not None and max_length is None:
        if truncation:
         UpperCAmelCase__							 =						min(array.shape[0] for array in raw_audio      )
         UpperCAmelCase__							 =						int(np.floor(max_length / self.chunk_stride      )      )
         UpperCAmelCase__							 =						(nb_step - 1) * self.chunk_stride + self.chunk_length
        elif padding:
         UpperCAmelCase__							 =						max(array.shape[0] for array in raw_audio      )
         UpperCAmelCase__							 =						int(np.ceil(max_length / self.chunk_stride      )      )
         UpperCAmelCase__							 =						(nb_step - 1) * self.chunk_stride + self.chunk_length
         UpperCAmelCase__							 =						"max_length"
        else:
         UpperCAmelCase__							 =						input_values

        # normal padding on batch
       if padded_inputs is None:
        UpperCAmelCase__							 =						self.pad(
            lowerCamelCase ,      max_length=lowerCamelCase ,      truncation=lowerCamelCase ,      padding=lowerCamelCase ,      return_attention_mask=lowerCamelCase ,      )
        if padding:
         UpperCAmelCase__							 =						padded_inputs.pop("attention_mask"      )

       UpperCAmelCase__							 =						[]
       for example in padded_inputs.pop("input_values"      ):
        if self.feature_size == 1:
         UpperCAmelCase__							 =						example[..., None]
        input_values.append(example.T      )

       UpperCAmelCase__							 =						input_values
       if return_tensors is not None:
        UpperCAmelCase__							 =						padded_inputs.convert_to_tensors(lowerCamelCase      )

       return padded_inputs





