


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 						lowerCamelCase_							:



					'''simple docstring'''







					def __init__(					self			,							snake_case_			,							snake_case_=1_3			,							snake_case_=7			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=True			,							snake_case_=9_9			,							snake_case_=3_2			,							snake_case_=5			,							snake_case_=4			,							snake_case_=3_7			,							snake_case_="gelu"			,							snake_case_=0.1			,							snake_case_=0.1			,							snake_case_=5_1_2			,							snake_case_=1_6			,							snake_case_=2			,							snake_case_=0.0_2			,							snake_case_=3			,							snake_case_=4			,							snake_case_=None			,							) ->      Optional[Any]:







								'''simple docstring'''

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






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

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

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

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

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

								__lowercase			       = self.get_config()

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






					def A   (					self    ) ->      List[Any]:







								'''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=snake_case_			,							initializer_range=self.initializer_range			,							)






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      int:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Optional[int]:







								'''simple docstring'''

								__lowercase			       = NystromformerForQuestionAnswering(config=snake_case_    )
								model.to(snake_case_    )
								model.eval()
								__lowercase			       = model(
								    snake_case_			,							attention_mask=snake_case_			,							token_type_ids=snake_case_			,							start_positions=snake_case_			,							end_positions=snake_case_			,							)
								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 A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Union[str, Any]:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      Tuple:







								'''simple docstring'''

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






					def A   (					self			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_			,							snake_case_    ) ->      int:







								'''simple docstring'''

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






					def A   (					self    ) ->      Any:







								'''simple docstring'''

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




@require_torch
class 						lowerCamelCase_							(					lowerCAmelCase__  ,  lowerCAmelCase__  ,  unittest.TestCase					):



					'''simple docstring'''







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






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







								'''simple docstring'''

								__lowercase			       = NystromformerModelTester(self    )
								__lowercase			       = ConfigTester(self			,							config_class=snake_case_			,							hidden_size=3_7    )






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







								'''simple docstring'''

								self.config_tester.run_common_tests()






					def A   (					self    ) ->      int:







								'''simple docstring'''

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






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







								'''simple docstring'''

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






					def A   (					self    ) ->      str:







								'''simple docstring'''

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






					def A   (					self    ) ->      Dict:







								'''simple docstring'''

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






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







								'''simple docstring'''

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






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







								'''simple docstring'''

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






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







								'''simple docstring'''

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






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







								'''simple docstring'''

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




@require_torch
class 						lowerCamelCase_							(					unittest.TestCase					):



					'''simple docstring'''







					@slow
					def A   (					self    ) ->      Dict:







								'''simple docstring'''

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

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

								__lowercase			       = torch.Size((1, 6, 7_6_8)    )
								self.assertEqual(output.shape			,							snake_case_    )

								__lowercase			       = torch.tensor(
								    [[[-0.4_5_3_2, -0.0_9_3_6, 0.5_1_3_7], [-0.2_6_7_6, 0.0_6_2_8, 0.6_1_8_6], [-0.3_6_2_9, -0.1_7_2_6, 0.4_7_1_6]]]    )

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






					@slow
					def A   (					self    ) ->      List[str]:







								'''simple docstring'''

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

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

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

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

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

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






import copy
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Optional, Union
@dataclass
class 						lowerCamelCase_							:



					'''simple docstring'''







					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					True
					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					1
					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					False
					__UpperCAmelCase							  =					None
					__UpperCAmelCase							  =					None






					def A   (					self    ) ->      "DownloadConfig":







								'''simple docstring'''

								return self.__class__(**{k: copy.deepcopy(snake_case_    ) for k, v in self.__dict__.items()}    )


