

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  :






			"""simple docstring"""




			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase=13 ,				__UpperCamelCase=7 ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=99 ,				__UpperCamelCase=32 ,				__UpperCamelCase=5 ,				__UpperCamelCase=4 ,				__UpperCamelCase=37 ,				__UpperCamelCase="gelu" ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=0.1 ,				__UpperCamelCase=5_12 ,				__UpperCamelCase=16 ,				__UpperCamelCase=2 ,				__UpperCamelCase=0.02 ,				__UpperCamelCase=3 ,				__UpperCamelCase=4 ,				__UpperCamelCase=None ,				):

				"""simple docstring"""


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


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size							)

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

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

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

				snake_case_      		= self.get_config()

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


			def 			__lowerCAmelCase (    self							):

				"""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=__UpperCamelCase ,				initializer_range=self.initializer_range ,				)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= NystromformerModel(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= NystromformerForMaskedLM(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.vocab_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= NystromformerForQuestionAnswering(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				start_positions=__UpperCamelCase ,				end_positions=__UpperCamelCase ,				)
				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 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= NystromformerForSequenceClassification(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= NystromformerForTokenClassification(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


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







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.prepare_config_and_inputs()
				(
				    (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       
				)      		= config_and_inputs
				snake_case_      		= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
				return config, inputs_dict







@require_torch
class SCREAMING_SNAKE_CASE  (		__snake_case							,						__snake_case							,						unittest.TestCase							):






			"""simple docstring"""




			__A						  =  (
			    (
			        NystromformerModel,
			        NystromformerForMaskedLM,
			        NystromformerForMultipleChoice,
			        NystromformerForQuestionAnswering,
			        NystromformerForSequenceClassification,
			        NystromformerForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__A						  =  (
			    {
			        """feature-extraction""": NystromformerModel,
			        """fill-mask""": NystromformerForMaskedLM,
			        """question-answering""": NystromformerForQuestionAnswering,
			        """text-classification""": NystromformerForSequenceClassification,
			        """token-classification""": NystromformerForTokenClassification,
			        """zero-shot""": NystromformerForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			__A						  =  False
			__A						  =  False


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= NystromformerModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase ,				hidden_size=37							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
					snake_case_      		= type
					self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase							)







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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







@require_torch
class SCREAMING_SNAKE_CASE  (		unittest.TestCase							):






			"""simple docstring"""




			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


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

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= torch.Size((1, 6, 7_68)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= 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] ,				__UpperCamelCase ,				atol=1E-4							)							)







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= 'the [MASK] of Belgium is Brussels'

				snake_case_      		= AutoTokenizer.from_pretrained('uw-madison/nystromformer-512'							)
				snake_case_      		= NystromformerForMaskedLM.from_pretrained('uw-madison/nystromformer-512'							)

				snake_case_      		= tokenizer(__UpperCamelCase ,				return_tensors='pt'							)

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

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

				self.assertEqual(tokenizer.decode(__UpperCamelCase							) ,				'capital'							)



from ..utils import DummyObject, requires_backends


class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







class SCREAMING_SNAKE_CASE  (		metaclass=__snake_case							):






			"""simple docstring"""




			__A						  =  ["""torch""", """transformers""", """onnx"""]


			def __init__(    self ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(self ,				['torch', 'transformers', 'onnx']							)


			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)







			@classmethod
			def 			__lowerCAmelCase (    cls ,				*__UpperCamelCase ,				**__UpperCamelCase							):

				"""simple docstring"""


				requires_backends(cls ,				['torch', 'transformers', 'onnx']							)
