



import unittest

from transformers import AutoTokenizer, NystromformerConfig, is_torch_available
from transformers.testing_utils import require_torch, 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 (
          NystromformerForMaskedLM,
          NystromformerForMultipleChoice,
          NystromformerForQuestionAnswering,
          NystromformerForSequenceClassification,
          NystromformerForTokenClassification,
          NystromformerModel,
      )
      from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST






class    SCREAMING_SNAKE_CASE			:





       def __init__(   self		:		Union[str, Any]      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		Any=13      ,				__lowercase		:		str=7      ,				__lowercase		:		Optional[Any]=True      ,				__lowercase		:		Tuple=True      ,				__lowercase		:		Dict=True      ,				__lowercase		:		int=True      ,				__lowercase		:		Optional[int]=99      ,				__lowercase		:		int=32      ,				__lowercase		:		List[str]=5      ,				__lowercase		:		int=4      ,				__lowercase		:		Tuple=37      ,				__lowercase		:		Dict="gelu"      ,				__lowercase		:		Tuple=0.1      ,				__lowercase		:		List[Any]=0.1      ,				__lowercase		:		str=512      ,				__lowercase		:		Optional[int]=16      ,				__lowercase		:		Any=2      ,				__lowercase		:		Any=0.02      ,				__lowercase		:		Dict=3      ,				__lowercase		:		Optional[Any]=4      ,				__lowercase		:		Dict=None      ,				):

           '''simple docstring'''
           __a       =		parent
           __a       =		batch_size
           __a       =		seq_length
           __a       =		is_training
           __a       =		use_input_mask
           __a       =		use_token_type_ids
           __a       =		use_labels
           __a       =		vocab_size
           __a       =		hidden_size
           __a       =		num_hidden_layers
           __a       =		num_attention_heads
           __a       =		intermediate_size
           __a       =		hidden_act
           __a       =		hidden_dropout_prob
           __a       =		attention_probs_dropout_prob
           __a       =		max_position_embeddings
           __a       =		type_vocab_size
           __a       =		type_sequence_label_size
           __a       =		initializer_range
           __a       =		num_labels
           __a       =		num_choices
           __a       =		scope





       def 						UpperCamelCase_					(   self		:		Dict	):

           '''simple docstring'''
           __a       =		ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size	)

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

           __a       =		None
           if self.use_token_type_ids:
               __a       =		ids_tensor([self.batch_size, self.seq_length]      ,				self.type_vocab_size	)

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

           __a       =		self.get_config()

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





       def 						UpperCamelCase_					(   self		:		List[str]	):

           '''simple docstring'''
           return NystromformerConfig(
               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=a_      ,				initializer_range=self.initializer_range      ,				)





       def 						UpperCamelCase_					(   self		:		Tuple      ,				__lowercase		:		List[Any]      ,				__lowercase		:		Any      ,				__lowercase		:		Dict      ,				__lowercase		:		Dict      ,				__lowercase		:		Any      ,				__lowercase		:		Dict      ,				__lowercase		:		Tuple	):

           '''simple docstring'''
           __a       =		NystromformerModel(config=a_	)
           model.to(a_	)
           model.eval()
           __a       =		model(a_      ,				attention_mask=a_      ,				token_type_ids=a_	)
           __a       =		model(a_      ,				token_type_ids=a_	)
           __a       =		model(a_	)
           self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)	)





       def 						UpperCamelCase_					(   self		:		Any      ,				__lowercase		:		int      ,				__lowercase		:		str      ,				__lowercase		:		Any      ,				__lowercase		:		List[Any]      ,				__lowercase		:		Union[str, Any]      ,				__lowercase		:		Union[str, Any]      ,				__lowercase		:		str	):

           '''simple docstring'''
           __a       =		NystromformerForMaskedLM(config=a_	)
           model.to(a_	)
           model.eval()
           __a       =		model(a_      ,				attention_mask=a_      ,				token_type_ids=a_      ,				labels=a_	)
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)	)





       def 						UpperCamelCase_					(   self		:		Optional[Any]      ,				__lowercase		:		List[Any]      ,				__lowercase		:		Any      ,				__lowercase		:		int      ,				__lowercase		:		Union[str, Any]      ,				__lowercase		:		List[str]      ,				__lowercase		:		Any      ,				__lowercase		:		Tuple	):

           '''simple docstring'''
           __a       =		NystromformerForQuestionAnswering(config=a_	)
           model.to(a_	)
           model.eval()
           __a       =		model(
               a_      ,				attention_mask=a_      ,				token_type_ids=a_      ,				start_positions=a_      ,				end_positions=a_      ,				)
           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		:		Union[str, Any]      ,				__lowercase		:		Any      ,				__lowercase		:		Optional[int]      ,				__lowercase		:		Tuple      ,				__lowercase		:		List[Any]      ,				__lowercase		:		Dict      ,				__lowercase		:		Optional[int]      ,				__lowercase		:		Any	):

           '''simple docstring'''
           __a       =		self.num_labels
           __a       =		NystromformerForSequenceClassification(a_	)
           model.to(a_	)
           model.eval()
           __a       =		model(a_      ,				attention_mask=a_      ,				token_type_ids=a_      ,				labels=a_	)
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_labels)	)





       def 						UpperCamelCase_					(   self		:		Optional[Any]      ,				__lowercase		:		str      ,				__lowercase		:		Dict      ,				__lowercase		:		Dict      ,				__lowercase		:		List[Any]      ,				__lowercase		:		List[Any]      ,				__lowercase		:		str      ,				__lowercase		:		Union[str, Any]	):

           '''simple docstring'''
           __a       =		self.num_labels
           __a       =		NystromformerForTokenClassification(config=a_	)
           model.to(a_	)
           model.eval()
           __a       =		model(a_      ,				attention_mask=a_      ,				token_type_ids=a_      ,				labels=a_	)
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.num_labels)	)





       def 						UpperCamelCase_					(   self		:		Optional[Any]      ,				__lowercase		:		Tuple      ,				__lowercase		:		str      ,				__lowercase		:		Optional[int]      ,				__lowercase		:		Any      ,				__lowercase		:		Optional[Any]      ,				__lowercase		:		int      ,				__lowercase		:		Dict	):

           '''simple docstring'''
           __a       =		self.num_choices
           __a       =		NystromformerForMultipleChoice(config=a_	)
           model.to(a_	)
           model.eval()
           __a       =		input_ids.unsqueeze(1	).expand(-1      ,				self.num_choices      ,				-1	).contiguous()
           __a       =		token_type_ids.unsqueeze(1	).expand(-1      ,				self.num_choices      ,				-1	).contiguous()
           __a       =		input_mask.unsqueeze(1	).expand(-1      ,				self.num_choices      ,				-1	).contiguous()
           __a       =		model(
               a_      ,				attention_mask=a_      ,				token_type_ids=a_      ,				labels=a_      ,				)
           self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_choices)	)







       def 						UpperCamelCase_					(   self		:		str	):

           '''simple docstring'''
           __a       =		self.prepare_config_and_inputs()
           (
               __a      
           )       =		config_and_inputs
           __a       =		{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
           return config, inputs_dict






@require_torch
class    SCREAMING_SNAKE_CASE			(							__a			,      __a			,      unittest.TestCase					):
       __lowerCamelCase	:	int    =(
           (
               NystromformerModel,
               NystromformerForMaskedLM,
               NystromformerForMultipleChoice,
               NystromformerForQuestionAnswering,
               NystromformerForSequenceClassification,
               NystromformerForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       __lowerCamelCase	:	Optional[Any]    =(
           {
               'feature-extraction': NystromformerModel,
               'fill-mask': NystromformerForMaskedLM,
               'question-answering': NystromformerForQuestionAnswering,
               'text-classification': NystromformerForSequenceClassification,
               'token-classification': NystromformerForTokenClassification,
               'zero-shot': NystromformerForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       __lowerCamelCase	:	Union[str, Any]    =False
       __lowerCamelCase	:	Tuple    =False





       def 						UpperCamelCase_					(   self		:		str	):

           '''simple docstring'''
           __a       =		NystromformerModelTester(self	)
           __a       =		ConfigTester(self      ,				config_class=a_      ,				hidden_size=37	)





       def 						UpperCamelCase_					(   self		:		List[str]	):

           '''simple docstring'''
           self.config_tester.run_common_tests()





       def 						UpperCamelCase_					(   self		:		int	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*a_	)





       def 						UpperCamelCase_					(   self		:		Dict	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
               __a       =		type
               self.model_tester.create_and_check_model(*a_	)





       def 						UpperCamelCase_					(   self		:		Optional[Any]	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*a_	)





       def 						UpperCamelCase_					(   self		:		Dict	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*a_	)





       def 						UpperCamelCase_					(   self		:		Optional[Any]	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*a_	)





       def 						UpperCamelCase_					(   self		:		Dict	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*a_	)





       def 						UpperCamelCase_					(   self		:		str	):

           '''simple docstring'''
           __a       =		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*a_	)







       @slow
       def 						UpperCamelCase_					(   self		:		int	):

           '''simple docstring'''



           for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               __a       =		NystromformerModel.from_pretrained(a_	)
               self.assertIsNotNone(a_	)






@require_torch
class    SCREAMING_SNAKE_CASE			(							unittest.TestCase					):





       @slow
       def 						UpperCamelCase_					(   self		:		List[Any]	):

           '''simple docstring'''
           __a       =		NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""	)
           __a       =		torch.tensor([[0, 1, 2, 3, 4, 5]]	)

           with torch.no_grad():
               __a       =		model(a_	)[0]

           __a       =		torch.Size((1, 6, 768)	)
           self.assertEqual(output.shape      ,				a_	)

           __a       =		torch.tensor(
               [[[-0.4532, -0.0936, 0.5137], [-0.2676, 0.0628, 0.6186], [-0.3629, -0.1726, 0.4716]]]	)

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







       @slow
       def 						UpperCamelCase_					(   self		:		Any	):

           '''simple docstring'''
           __a       =		"""the [MASK] of Belgium is Brussels"""

           __a       =		AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""	)
           __a       =		NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""	)

           __a       =		tokenizer(a_      ,				return_tensors="""pt"""	)

           with torch.no_grad():
               __a       =		model(encoding.input_ids	).logits

           __a       =		token_logits[:, 2, :].argmax(-1	)[0]

           self.assertEqual(tokenizer.decode(a_	)      ,				"""capital"""	)





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


lowerCamelCase__			= {"""configuration_speech_encoder_decoder""": ["""SpeechEncoderDecoderConfig"""]}

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""SpeechEncoderDecoderModel"""]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""FlaxSpeechEncoderDecoderModel"""]

if TYPE_CHECKING:
						from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
						import sys

						lowerCamelCase__			= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
