

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

				def __init__(			self :      List[str]							,    snake_case__ :      Dict							,    snake_case__ :      Optional[Any]=1_3							,    snake_case__ :      Union[str, Any]=7							,    snake_case__ :      List[str]=True							,    snake_case__ :      Any=True							,    snake_case__ :      List[str]=True							,    snake_case__ :      Optional[Any]=True							,    snake_case__ :      List[str]=9_9							,    snake_case__ :      str=3_2							,    snake_case__ :      Dict=5							,    snake_case__ :      str=4							,    snake_case__ :      int=3_7							,    snake_case__ :      Union[str, Any]="gelu"							,    snake_case__ :      Any=0.1							,    snake_case__ :      Dict=0.1							,    snake_case__ :      Optional[Any]=5_1_2							,    snake_case__ :      List[Any]=1_6							,    snake_case__ :      str=2							,    snake_case__ :      int=0.02							,    snake_case__ :      List[str]=3							,    snake_case__ :      Dict=4							,    snake_case__ :      str=None							,    ):

					"""simple docstring"""






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

				def UpperCamelCase			(			self :      List[str]    ):

					"""simple docstring"""






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

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

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

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

					SCREAMING_SNAKE_CASE								=			self.get_config()

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

				def UpperCamelCase			(			self :      Optional[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 UpperCamelCase			(			self :      List[str]							,    snake_case__ :      Any							,    snake_case__ :      Tuple							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Dict							,    snake_case__ :      List[str]							,    snake_case__ :      int							,    snake_case__ :      Union[str, Any]    ):

					"""simple docstring"""






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

				def UpperCamelCase			(			self :      List[Any]							,    snake_case__ :      str							,    snake_case__ :      Dict							,    snake_case__ :      Dict							,    snake_case__ :      Dict							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      Tuple    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			NystromformerForMaskedLM(config=snake_case__    )
					model.to(snake_case__    )
					model.eval()
					SCREAMING_SNAKE_CASE								=			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 UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      Tuple							,    snake_case__ :      Optional[Any]							,    snake_case__ :      int							,    snake_case__ :      Optional[int]							,    snake_case__ :      str							,    snake_case__ :      List[str]							,    snake_case__ :      Optional[Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			NystromformerForQuestionAnswering(config=snake_case__    )
					model.to(snake_case__    )
					model.eval()
					SCREAMING_SNAKE_CASE								=			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 UpperCamelCase			(			self :      List[str]							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      str							,    snake_case__ :      Optional[int]							,    snake_case__ :      Any							,    snake_case__ :      Union[str, Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.num_labels
					SCREAMING_SNAKE_CASE								=			NystromformerForSequenceClassification(snake_case__    )
					model.to(snake_case__    )
					model.eval()
					SCREAMING_SNAKE_CASE								=			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 UpperCamelCase			(			self :      Tuple							,    snake_case__ :      int							,    snake_case__ :      Tuple							,    snake_case__ :      List[str]							,    snake_case__ :      List[Any]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Dict    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.num_labels
					SCREAMING_SNAKE_CASE								=			NystromformerForTokenClassification(config=snake_case__    )
					model.to(snake_case__    )
					model.eval()
					SCREAMING_SNAKE_CASE								=			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 UpperCamelCase			(			self :      int							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[str]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Tuple							,    snake_case__ :      Optional[int]							,    snake_case__ :      Any							,    snake_case__ :      List[Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.num_choices
					SCREAMING_SNAKE_CASE								=			NystromformerForMultipleChoice(config=snake_case__    )
					model.to(snake_case__    )
					model.eval()
					SCREAMING_SNAKE_CASE								=			input_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
					SCREAMING_SNAKE_CASE								=			token_type_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
					SCREAMING_SNAKE_CASE								=			input_mask.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
					SCREAMING_SNAKE_CASE								=			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 UpperCamelCase			(			self :      Union[str, Any]    ):

					"""simple docstring"""






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




@require_torch
class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
				__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 UpperCamelCase			(			self :      int    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			NystromformerModelTester(self    )
					SCREAMING_SNAKE_CASE								=			ConfigTester(self							,    config_class=snake_case__							,    hidden_size=3_7    )

				def UpperCamelCase			(			self :      List[str]    ):

					"""simple docstring"""






					self.config_tester.run_common_tests()

				def UpperCamelCase			(			self :      List[Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*snake_case__    )

				def UpperCamelCase			(			self :      int    ):

					"""simple docstring"""






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

				def UpperCamelCase			(			self :      List[Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*snake_case__    )

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

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_multiple_choice(*snake_case__    )

				def UpperCamelCase			(			self :      Tuple    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_question_answering(*snake_case__    )

				def UpperCamelCase			(			self :      Any    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_sequence_classification(*snake_case__    )

				def UpperCamelCase			(			self :      List[Any]    ):

					"""simple docstring"""






					SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*snake_case__    )




				@slow
				def UpperCamelCase			(			self :      List[str]    ):

					"""simple docstring"""






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




@require_torch
class      UpperCamelCase (     unittest.TestCase   ):

				@slow
				def UpperCamelCase			(			self :      int    ):

					"""simple docstring"""






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

					with torch.no_grad():
						SCREAMING_SNAKE_CASE								=			model(snake_case__    )[0]

					SCREAMING_SNAKE_CASE								=			torch.Size((1, 6, 7_6_8)    )
					self.assertEqual(output.shape							,    snake_case__    )

					SCREAMING_SNAKE_CASE								=			torch.tensor(
					    [[[-0.4_532, -0.0_936, 0.5_137], [-0.2_676, 0.0_628, 0.6_186], [-0.3_629, -0.1_726, 0.4_716]]]    )

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




				@slow
				def UpperCamelCase			(			self :      Dict    ):

					"""simple docstring"""






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

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

					SCREAMING_SNAKE_CASE								=			tokenizer(snake_case__							,    return_tensors='pt'    )

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

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

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




from ....configuration_utils import PretrainedConfig
from ....utils import logging


a_    :						Any    		=      logging.get_logger(__name__)

a_    :						Dict    		=      {
    "Visual-Attention-Network/van-base": (
        "https://huggingface.co/Visual-Attention-Network/van-base/blob/main/config.json"
    ),
}




class      UpperCamelCase (     SCREAMING_SNAKE_CASE   ):
    __UpperCamelCase         ="van"

    def __init__(			self :      Optional[Any]							,    snake_case__ :      Tuple=2_2_4							,    snake_case__ :      Dict=3							,    snake_case__ :      Union[str, Any]=[7, 3, 3, 3]							,    snake_case__ :      str=[4, 2, 2, 2]							,    snake_case__ :      Optional[Any]=[6_4, 1_2_8, 3_2_0, 5_1_2]							,    snake_case__ :      Optional[Any]=[3, 3, 1_2, 3]							,    snake_case__ :      Tuple=[8, 8, 4, 4]							,    snake_case__ :      Any="gelu"							,    snake_case__ :      Dict=0.02							,    snake_case__ :      List[str]=1E-6							,    snake_case__ :      int=1E-2							,    snake_case__ :      Any=0.0							,    snake_case__ :      Tuple=0.0							,    **snake_case__ :      Any							,    ):

     """simple docstring"""






     super().__init__(**snake_case__    )
     SCREAMING_SNAKE_CASE								=			image_size
     SCREAMING_SNAKE_CASE								=			num_channels
     SCREAMING_SNAKE_CASE								=			patch_sizes
     SCREAMING_SNAKE_CASE								=			strides
     SCREAMING_SNAKE_CASE								=			hidden_sizes
     SCREAMING_SNAKE_CASE								=			depths
     SCREAMING_SNAKE_CASE								=			mlp_ratios
     SCREAMING_SNAKE_CASE								=			hidden_act
     SCREAMING_SNAKE_CASE								=			initializer_range
     SCREAMING_SNAKE_CASE								=			layer_norm_eps
     SCREAMING_SNAKE_CASE								=			layer_scale_init_value
     SCREAMING_SNAKE_CASE								=			drop_path_rate
     SCREAMING_SNAKE_CASE								=			dropout_rate

