
import os
import tempfile
import unittest

from transformers import NezhaConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_gpu, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
    import torch

    from transformers import (
        MODEL_FOR_PRETRAINING_MAPPING,
        NezhaForMaskedLM,
        NezhaForMultipleChoice,
        NezhaForNextSentencePrediction,
        NezhaForPreTraining,
        NezhaForQuestionAnswering,
        NezhaForSequenceClassification,
        NezhaForTokenClassification,
        NezhaModel,
    )
    from transformers.models.nezha.modeling_nezha import NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST




class 	__A:

  """simple docstring"""







  def __init__(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_=13       ,			SCREAMING_SNAKE_CASE_=7       ,			SCREAMING_SNAKE_CASE_=True       ,			SCREAMING_SNAKE_CASE_=True       ,			SCREAMING_SNAKE_CASE_=True       ,			SCREAMING_SNAKE_CASE_=True       ,			SCREAMING_SNAKE_CASE_=99       ,			SCREAMING_SNAKE_CASE_=32       ,			SCREAMING_SNAKE_CASE_=5       ,			SCREAMING_SNAKE_CASE_=4       ,			SCREAMING_SNAKE_CASE_=37       ,			SCREAMING_SNAKE_CASE_="gelu"       ,			SCREAMING_SNAKE_CASE_=0.1       ,			SCREAMING_SNAKE_CASE_=0.1       ,			SCREAMING_SNAKE_CASE_=1_28       ,			SCREAMING_SNAKE_CASE_=32       ,			SCREAMING_SNAKE_CASE_=16       ,			SCREAMING_SNAKE_CASE_=2       ,			SCREAMING_SNAKE_CASE_=0.02       ,			SCREAMING_SNAKE_CASE_=3       ,			SCREAMING_SNAKE_CASE_=4       ,			SCREAMING_SNAKE_CASE_=None       ,			):
       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						):
       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
       if self.use_token_type_ids:
            UpperCamelCase__				=				ids_tensor([self.batch_size, self.seq_length]       ,			self.type_vocab_size						)

       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, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







  def        UpperCAmelCase_					(self						):
       return NezhaConfig(
           vocab_size=self.vocab_size       ,			hidden_size=self.hidden_size       ,			num_hidden_layers=self.num_hidden_layers       ,			num_attention_heads=self.num_attention_heads       ,			intermediate_size=self.intermediate_size       ,			hidden_act=self.hidden_act       ,			hidden_dropout_prob=self.hidden_dropout_prob       ,			attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,			max_position_embeddings=self.max_position_embeddings       ,			type_vocab_size=self.type_vocab_size       ,			is_decoder=_lowerCamelCase       ,			initializer_range=self.initializer_range       ,			)







  def        UpperCAmelCase_					(self						):
       (
           (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,           
       )				=				self.prepare_config_and_inputs()

       UpperCamelCase__				=				True
       UpperCamelCase__				=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
       UpperCamelCase__				=				ids_tensor([self.batch_size, self.seq_length]       ,			vocab_size=2						)

       return (
           config,
           input_ids,
           token_type_ids,
           input_mask,
           sequence_labels,
           token_labels,
           choice_labels,
           encoder_hidden_states,
           encoder_attention_mask,
       )







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				NezhaModel(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase						)
       UpperCamelCase__				=				model(_lowerCamelCase       ,			token_type_ids=_lowerCamelCase						)
       UpperCamelCase__				=				model(_lowerCamelCase						)
       self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size)						)
       self.parent.assertEqual(result.pooler_output.shape       ,			(self.batch_size, self.hidden_size)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			):
       UpperCamelCase__				=				True
       UpperCamelCase__				=				NezhaModel(_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			encoder_hidden_states=_lowerCamelCase       ,			encoder_attention_mask=_lowerCamelCase       ,			)
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			encoder_hidden_states=_lowerCamelCase       ,			)
       UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase						)
       self.parent.assertEqual(result.last_hidden_state.shape       ,			(self.batch_size, self.seq_length, self.hidden_size)						)
       self.parent.assertEqual(result.pooler_output.shape       ,			(self.batch_size, self.hidden_size)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				NezhaForMaskedLM(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				NezhaForNextSentencePrediction(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase       ,			)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, 2)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				NezhaForPreTraining(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase       ,			next_sentence_label=_lowerCamelCase       ,			)
       self.parent.assertEqual(result.prediction_logits.shape       ,			(self.batch_size, self.seq_length, self.vocab_size)						)
       self.parent.assertEqual(result.seq_relationship_logits.shape       ,			(self.batch_size, 2)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				NezhaForQuestionAnswering(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			start_positions=_lowerCamelCase       ,			end_positions=_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        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				self.num_labels
       UpperCamelCase__				=				NezhaForSequenceClassification(_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_labels)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				self.num_labels
       UpperCamelCase__				=				NezhaForTokenClassification(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase						)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.seq_length, self.num_labels)						)







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_						):
       UpperCamelCase__				=				self.num_choices
       UpperCamelCase__				=				NezhaForMultipleChoice(config=_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()
       UpperCamelCase__				=				input_ids.unsqueeze(1						).expand(-1       ,			self.num_choices       ,			-1						).contiguous()
       UpperCamelCase__				=				token_type_ids.unsqueeze(1						).expand(-1       ,			self.num_choices       ,			-1						).contiguous()
       UpperCamelCase__				=				input_mask.unsqueeze(1						).expand(-1       ,			self.num_choices       ,			-1						).contiguous()
       UpperCamelCase__				=				model(
           _lowerCamelCase       ,			attention_mask=_lowerCamelCase       ,			token_type_ids=_lowerCamelCase       ,			labels=_lowerCamelCase       ,			)
       self.parent.assertEqual(result.logits.shape       ,			(self.batch_size, self.num_choices)						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,           
       )				=				config_and_inputs
       UpperCamelCase__				=				{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
       return config, inputs_dict






@require_torch
class 	__A(				_lowerCAmelCase				,      _lowerCAmelCase				,      _lowerCAmelCase				,      unittest.TestCase ):

  """simple docstring"""
  SCREAMING_SNAKE_CASE__			    =			(
      (
          NezhaModel,
          NezhaForMaskedLM,
          NezhaForMultipleChoice,
          NezhaForNextSentencePrediction,
          NezhaForPreTraining,
          NezhaForQuestionAnswering,
          NezhaForSequenceClassification,
          NezhaForTokenClassification,
      )
      if is_torch_available()
      else ()
  )
  SCREAMING_SNAKE_CASE__			    =			(
      {
          "feature-extraction": NezhaModel,
          "fill-mask": NezhaForMaskedLM,
          "question-answering": NezhaForQuestionAnswering,
          "text-classification": NezhaForSequenceClassification,
          "token-classification": NezhaForTokenClassification,
          "zero-shot": NezhaForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  SCREAMING_SNAKE_CASE__			    =			True







  def        UpperCAmelCase_					(self       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_       ,			SCREAMING_SNAKE_CASE_=False						):
       UpperCamelCase__				=				super()._prepare_for_class(_lowerCamelCase       ,			_lowerCamelCase       ,			return_labels=_lowerCamelCase						)

       if return_labels:
            if model_class in get_values(_lowerCamelCase						):
                 UpperCamelCase__				=				torch.zeros(
                     (self.model_tester.batch_size, self.model_tester.seq_length)       ,			dtype=torch.long       ,			device=_lowerCamelCase						)
                 UpperCamelCase__				=				torch.zeros(
                     self.model_tester.batch_size       ,			dtype=torch.long       ,			device=_lowerCamelCase						)
       return inputs_dict







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				NezhaModelTester(self						)
       UpperCamelCase__				=				ConfigTester(self       ,			config_class=_lowerCamelCase       ,			hidden_size=37						)







  def        UpperCAmelCase_					(self						):
       self.config_tester.run_common_tests()







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       (
           (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,       (
           UpperCamelCase__    
       )		,           
       )				=				self.model_tester.prepare_config_and_inputs_for_decoder()

       UpperCamelCase__				=				None

       self.model_tester.create_and_check_model_as_decoder(
           _lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			_lowerCamelCase       ,			)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_next_sequence_prediction(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_pretraining(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase						)







  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase						)







  @slow
  def        UpperCAmelCase_					(self						):
       for model_name in NEZHA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCamelCase__				=				NezhaModel.from_pretrained(_lowerCamelCase						)
            self.assertIsNotNone(_lowerCamelCase						)







  @slow
  @require_torch_gpu
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__		,       UpperCamelCase__				=				self.model_tester.prepare_config_and_inputs_for_common()
       for model_class in self.all_model_classes:
            # NezhaForMultipleChoice behaves incorrectly in JIT environments.
            if model_class == NezhaForMultipleChoice:
                 return

            UpperCamelCase__				=				True
            UpperCamelCase__				=				model_class(config=_lowerCamelCase						)

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

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






@require_torch
class 	__A(				unittest.TestCase ):

  """simple docstring"""







  @slow
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				NezhaModel.from_pretrained("""sijunhe/nezha-cn-base"""						)
       UpperCamelCase__				=				torch.tensor([[0, 1, 2, 3, 4, 5]]						)
       UpperCamelCase__				=				torch.tensor([[0, 1, 1, 1, 1, 1]]						)
       with torch.no_grad():
            UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase						)[0]
       UpperCamelCase__				=				torch.Size((1, 6, 7_68)						)
       self.assertEqual(output.shape       ,			_lowerCamelCase						)
       UpperCamelCase__				=				torch.tensor([[[0.0685, 0.2441, 0.1102], [0.0600, 0.1906, 0.1349], [0.0221, 0.0819, 0.0586]]]						)

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







  @slow
  def        UpperCAmelCase_					(self						):
       UpperCamelCase__				=				NezhaForMaskedLM.from_pretrained("""sijunhe/nezha-cn-base"""						)
       UpperCamelCase__				=				torch.tensor([[0, 1, 2, 3, 4, 5]]						)
       UpperCamelCase__				=				torch.tensor([[1, 1, 1, 1, 1, 1]]						)
       with torch.no_grad():
            UpperCamelCase__				=				model(_lowerCamelCase       ,			attention_mask=_lowerCamelCase						)[0]
       UpperCamelCase__				=				torch.Size((1, 6, 2_11_28)						)
       self.assertEqual(output.shape       ,			_lowerCamelCase						)
       UpperCamelCase__				=				torch.tensor(
           [[-2.7939, -1.7902, -2.2189], [-2.8585, -1.8908, -2.3723], [-2.6499, -1.7750, -2.2558]]						)

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













"""simple docstring"""


import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


lowercase_				     =     random.Random()

if is_torch_available():
       import torch



def 					lowerCAmelCase				(    __UpperCamelCase		,       __UpperCamelCase=1.0		,       __UpperCamelCase=None		,       __UpperCamelCase=None ):




  """simple docstring"""







  if rng is None:
    __A									= global_rng

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

  return values


class 					snake_case		(							unittest.TestCase						):


   '''simple docstring'''

   def __init__(	self :     Any,       _lowerCamelCase :     List[str],       _lowerCamelCase :     Any=7,       _lowerCamelCase :     Optional[int]=4_00,       _lowerCamelCase :     Optional[int]=20_00,       _lowerCamelCase :     Dict=1,       _lowerCamelCase :     Optional[Any]=0.0,       _lowerCamelCase :     int=1_60_00,       _lowerCamelCase :     Optional[int]=True,       _lowerCamelCase :     Dict=True,       ):

     '''simple docstring'''
     __A									= parent
     __A									= batch_size
     __A									= min_seq_length
     __A									= max_seq_length
     __A									= (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
     __A									= feature_size
     __A									= padding_value
     __A									= sampling_rate
     __A									= return_attention_mask
     __A									= do_normalize


   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any]  ):

     '''simple docstring'''
     return {
         "feature_size": self.feature_size,
         "padding_value": self.padding_value,
         "sampling_rate": self.sampling_rate,
         "return_attention_mask": self.return_attention_mask,
         "do_normalize": self.do_normalize,
     }






   def 				_SCREAMING_SNAKE_CASE				(	self :     Any,       _lowerCamelCase :     Optional[Any]=False,       _lowerCamelCase :     int=False  ):

     '''simple docstring'''
     def _flatten(_lowerCamelCase :     List[str]  ):
       return list(itertools.chain(*_lowerCamelCase  )  )

     if equal_length:
       __A									= floats_list((self.batch_size, self.max_seq_length)  )
     else:
       # make sure that inputs increase in size
       __A									= [
           _flatten(floats_list((x, self.feature_size)  )  )
           for x in range(self.min_seq_length,       self.max_seq_length,       self.seq_length_diff  )
       ]

     if numpify:
       __A									= [np.asarray(_lowerCamelCase  ) for x in speech_inputs]

     return speech_inputs





@require_torch
@require_torchaudio
class 					snake_case		(							_lowerCAmelCase   ,	unittest.TestCase						):


   '''simple docstring'''

   A_						:	int						 = ASTFeatureExtractor


   def 				_SCREAMING_SNAKE_CASE				(	self :     Tuple  ):

     '''simple docstring'''
     __A									= ASTFeatureExtractionTester(self  )


   def 				_SCREAMING_SNAKE_CASE				(	self :     Dict  ):

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

     # Test not batched input
     __A									= feat_extract(speech_inputs[0],       return_tensors='''np'''  ).input_values
     __A									= feat_extract(np_speech_inputs[0],       return_tensors='''np'''  ).input_values
     self.assertTrue(np.allclose(_lowerCamelCase,       _lowerCamelCase,       atol=1e-3  )  )

     # Test batched
     __A									= feat_extract(_lowerCamelCase,       padding=_lowerCamelCase,       return_tensors='''np'''  ).input_values
     __A									= feat_extract(_lowerCamelCase,       padding=_lowerCamelCase,       return_tensors='''np'''  ).input_values
     for enc_seq_a, enc_seq_a in zip(_lowerCamelCase,       _lowerCamelCase  ):
       self.assertTrue(np.allclose(_lowerCamelCase,       _lowerCamelCase,       atol=1e-3  )  )

     # Test 2-D numpy arrays are batched.
     __A									= [floats_list((1, x)  )[0] for x in (8_00, 8_00, 8_00)]
     __A									= np.asarray(_lowerCamelCase  )
     __A									= feat_extract(_lowerCamelCase,       return_tensors='''np'''  ).input_values
     __A									= feat_extract(_lowerCamelCase,       return_tensors='''np'''  ).input_values
     for enc_seq_a, enc_seq_a in zip(_lowerCamelCase,       _lowerCamelCase  ):
       self.assertTrue(np.allclose(_lowerCamelCase,       _lowerCamelCase,       atol=1e-3  )  )


   @require_torch
   def 				_SCREAMING_SNAKE_CASE				(	self :     Union[str, Any]  ):

     '''simple docstring'''
     import torch

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

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


   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[int],       _lowerCamelCase :     Union[str, Any]  ):

     '''simple docstring'''
     from datasets import load_dataset

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

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






   @require_torch
   def 				_SCREAMING_SNAKE_CASE				(	self :     Optional[Any]  ):

     '''simple docstring'''
     # fmt: off
     __A									= torch.tensor(
         [-0.98_94, -1.27_76, -0.90_66, -1.27_76, -0.93_49, -1.26_09, -1.03_86, -1.27_76,
          -1.15_61, -1.27_76, -1.20_52, -1.27_23, -1.21_90, -1.21_32, -1.27_76, -1.11_33,
          -1.19_53, -1.13_43, -1.15_84, -1.22_03, -1.17_70, -1.24_74, -1.23_81, -1.19_36,
          -0.92_70, -0.83_17, -0.80_49, -0.77_06, -0.75_65, -0.78_69]  )
     # fmt: on

     __A									= self._load_datasamples(1  )
     __A									= ASTFeatureExtractor()
     __A									= feature_extractor(_lowerCamelCase,       return_tensors='''pt'''  ).input_values
     self.assertEquals(input_values.shape,       (1, 10_24, 1_28)  )
     self.assertTrue(torch.allclose(input_values[0, 0, :30],       _lowerCamelCase,       atol=1e-4  )  )




