




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



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[int]										=      parent
								_lowerCAmelCase    :    Dict										=      batch_size
								_lowerCAmelCase    :    int										=      seq_length
								_lowerCAmelCase    :    List[str]										=      is_training
								_lowerCAmelCase    :    str										=      use_input_mask
								_lowerCAmelCase    :    List[Any]										=      use_token_type_ids
								_lowerCAmelCase    :    Tuple										=      use_labels
								_lowerCAmelCase    :    Dict										=      vocab_size
								_lowerCAmelCase    :    Optional[int]										=      hidden_size
								_lowerCAmelCase    :    List[Any]										=      num_hidden_layers
								_lowerCAmelCase    :    List[Any]										=      num_attention_heads
								_lowerCAmelCase    :    Any										=      intermediate_size
								_lowerCAmelCase    :    Optional[int]										=      hidden_act
								_lowerCAmelCase    :    Optional[Any]										=      hidden_dropout_prob
								_lowerCAmelCase    :    List[str]										=      attention_probs_dropout_prob
								_lowerCAmelCase    :    Optional[int]										=      max_position_embeddings
								_lowerCAmelCase    :    Any										=      type_vocab_size
								_lowerCAmelCase    :    List[str]										=      type_sequence_label_size
								_lowerCAmelCase    :    Optional[Any]										=      initializer_range
								_lowerCAmelCase    :    Dict										=      num_labels
								_lowerCAmelCase    :    Optional[Any]										=      num_choices
								_lowerCAmelCase    :    List[Any]										=      scope



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Tuple:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size      )

								_lowerCAmelCase    :    List[Any]										=      None
								if self.use_input_mask:
												_lowerCAmelCase    :    Dict										=      random_attention_mask([self.batch_size, self.seq_length]      )

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

								_lowerCAmelCase    :    Dict										=      None
								_lowerCAmelCase    :    Optional[int]										=      None
								_lowerCAmelCase    :    Dict										=      None
								if self.use_labels:
												_lowerCAmelCase    :    Dict										=      ids_tensor([self.batch_size]		,	self.type_sequence_label_size      )
												_lowerCAmelCase    :    int										=      ids_tensor([self.batch_size, self.seq_length]		,	self.num_labels      )
												_lowerCAmelCase    :    Any										=      ids_tensor([self.batch_size]		,	self.num_choices      )

								_lowerCAmelCase    :    Any										=      self.get_config()

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



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    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=_UpperCAmelCase		,	initializer_range=self.initializer_range		,	)



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Union[str, Any]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Tuple		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		List[str]      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    Union[str, Any]										=      NystromformerModel(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Tuple										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	token_type_ids=_UpperCAmelCase      )
								_lowerCAmelCase    :    Tuple										=      model(_UpperCAmelCase		,	token_type_ids=_UpperCAmelCase      )
								_lowerCAmelCase    :    List[str]										=      model(_UpperCAmelCase      )
								self.parent.assertEqual(result.last_hidden_state.shape		,	(self.batch_size, self.seq_length, self.hidden_size)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]		,	_UpperCAmelCase      :		List[Any]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		Union[str, Any]      )  ->    int:





								'''simple docstring'''




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





								'''simple docstring'''




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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      self.num_labels
								_lowerCAmelCase    :    Optional[Any]										=      NystromformerForSequenceClassification(_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Tuple										=      model(_UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	token_type_ids=_UpperCAmelCase		,	labels=_UpperCAmelCase      )
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_labels)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]		,	_UpperCAmelCase      :		List[str]		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		Dict		,	_UpperCAmelCase      :		Any		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		str		,	_UpperCAmelCase      :		Dict      )  ->    Dict:





								'''simple docstring'''




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





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      self.num_choices
								_lowerCAmelCase    :    Union[str, Any]										=      NystromformerForMultipleChoice(config=_UpperCAmelCase      )
								model.to(_UpperCAmelCase      )
								model.eval()
								_lowerCAmelCase    :    Any										=      input_ids.unsqueeze(1      ).expand(-1		,	self.num_choices		,	-1      ).contiguous()
								_lowerCAmelCase    :    str										=      token_type_ids.unsqueeze(1      ).expand(-1		,	self.num_choices		,	-1      ).contiguous()
								_lowerCAmelCase    :    int										=      input_mask.unsqueeze(1      ).expand(-1		,	self.num_choices		,	-1      ).contiguous()
								_lowerCAmelCase    :    Union[str, Any]										=      model(
								    _UpperCAmelCase		,	attention_mask=_UpperCAmelCase		,	token_type_ids=_UpperCAmelCase		,	labels=_UpperCAmelCase		,	)
								self.parent.assertEqual(result.logits.shape		,	(self.batch_size, self.num_choices)      )



				def       SCREAMING_SNAKE_CASE       (      self      :		int      )  ->    Optional[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      self.prepare_config_and_inputs()
								(
								    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    (
								    _lowerCAmelCase
								)					,    
								)    :    Tuple										=      config_and_inputs
								_lowerCAmelCase    :    int										=      {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
								return config, inputs_dict






@require_torch
class        __snake_case  (_a						,	_a						,	unittest.TestCase		):
				lowerCAmelCase__				=     (
				    (
				        NystromformerModel,
				        NystromformerForMaskedLM,
				        NystromformerForMultipleChoice,
				        NystromformerForQuestionAnswering,
				        NystromformerForSequenceClassification,
				        NystromformerForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				lowerCAmelCase__				=     (
				    {
				        "feature-extraction": NystromformerModel,
				        "fill-mask": NystromformerForMaskedLM,
				        "question-answering": NystromformerForQuestionAnswering,
				        "text-classification": NystromformerForSequenceClassification,
				        "token-classification": NystromformerForTokenClassification,
				        "zero-shot": NystromformerForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				lowerCAmelCase__				=     False
				lowerCAmelCase__				=     False



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





								'''simple docstring'''




								_lowerCAmelCase    :    str										=      NystromformerModelTester(self      )
								_lowerCAmelCase    :    Optional[Any]										=      ConfigTester(self		,	config_class=_UpperCAmelCase		,	hidden_size=37      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    List[str]:





								'''simple docstring'''




								self.config_tester.run_common_tests()



				def       SCREAMING_SNAKE_CASE       (      self      :		str      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    Optional[int]:





								'''simple docstring'''




								_lowerCAmelCase    :    List[str]										=      self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
												_lowerCAmelCase    :    Dict										=      type
												self.model_tester.create_and_check_model(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    Any:





								'''simple docstring'''




								_lowerCAmelCase    :    Dict										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[Any]      )  ->    int:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		List[str]      )  ->    str:





								'''simple docstring'''




								_lowerCAmelCase    :    Tuple										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*_UpperCAmelCase      )



				def       SCREAMING_SNAKE_CASE       (      self      :		Tuple      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*_UpperCAmelCase      )



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





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase      )



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Optional[Any]      )  ->    List[str]:





								'''simple docstring'''




								for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												_lowerCAmelCase    :    List[str]										=      NystromformerModel.from_pretrained(_UpperCAmelCase      )
												self.assertIsNotNone(_UpperCAmelCase      )






@require_torch
class        __snake_case  (unittest.TestCase		):



				@slow
				def       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    List[Any]:





								'''simple docstring'''




								_lowerCAmelCase    :    Optional[Any]										=      NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""      )
								_lowerCAmelCase    :    int										=      torch.tensor([[0, 1, 2, 3, 4, 5]]      )

								with torch.no_grad():
												_lowerCAmelCase    :    Optional[Any]										=      model(_UpperCAmelCase      )[0]

								_lowerCAmelCase    :    Any										=      torch.Size((1, 6, 768)      )
								self.assertEqual(output.shape		,	_UpperCAmelCase      )

								_lowerCAmelCase    :    int										=      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       SCREAMING_SNAKE_CASE       (      self      :		Dict      )  ->    Dict:





								'''simple docstring'''




								_lowerCAmelCase    :    int										=      """the [MASK] of Belgium is Brussels"""

								_lowerCAmelCase    :    Union[str, Any]										=      AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""      )
								_lowerCAmelCase    :    Optional[Any]										=      NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""      )

								_lowerCAmelCase    :    Tuple										=      tokenizer(_UpperCAmelCase		,	return_tensors="""pt"""      )

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

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

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







import logging

from transformers.configuration_utils import PretrainedConfig


_lowerCamelCase      :  Union[str, Any]      						=						logging.getLogger(__name__)
class        __snake_case  (_a		):
				lowerCAmelCase__				=     "masked_bert"



				def __init__(      self      :		Union[str, Any]		,	_UpperCAmelCase      :		Union[str, Any]=3_0522		,	_UpperCAmelCase      :		Optional[Any]=768		,	_UpperCAmelCase      :		List[Any]=12		,	_UpperCAmelCase      :		Dict=12		,	_UpperCAmelCase      :		List[str]=3072		,	_UpperCAmelCase      :		List[Any]="gelu"		,	_UpperCAmelCase      :		Optional[Any]=0.1		,	_UpperCAmelCase      :		Optional[Any]=0.1		,	_UpperCAmelCase      :		List[Any]=512		,	_UpperCAmelCase      :		int=2		,	_UpperCAmelCase      :		Optional[int]=0.02		,	_UpperCAmelCase      :		str=1E-12		,	_UpperCAmelCase      :		Tuple=0		,	_UpperCAmelCase      :		Dict="topK"		,	_UpperCAmelCase      :		List[str]="constant"		,	_UpperCAmelCase      :		Optional[Any]=0.0		,	**_UpperCAmelCase      :		str		,	)  ->    Tuple:





								'''simple docstring'''




								super().__init__(pad_token_id=_UpperCAmelCase		,	**_UpperCAmelCase      )

								_lowerCAmelCase    :    str										=      vocab_size
								_lowerCAmelCase    :    List[str]										=      hidden_size
								_lowerCAmelCase    :    Any										=      num_hidden_layers
								_lowerCAmelCase    :    Optional[Any]										=      num_attention_heads
								_lowerCAmelCase    :    Optional[Any]										=      hidden_act
								_lowerCAmelCase    :    Union[str, Any]										=      intermediate_size
								_lowerCAmelCase    :    Tuple										=      hidden_dropout_prob
								_lowerCAmelCase    :    Optional[int]										=      attention_probs_dropout_prob
								_lowerCAmelCase    :    str										=      max_position_embeddings
								_lowerCAmelCase    :    int										=      type_vocab_size
								_lowerCAmelCase    :    List[str]										=      initializer_range
								_lowerCAmelCase    :    Optional[Any]										=      layer_norm_eps
								_lowerCAmelCase    :    Tuple										=      pruning_method
								_lowerCAmelCase    :    str										=      mask_init
								_lowerCAmelCase    :    List[str]										=      mask_scale

