

'''simple docstring'''





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    :		Optional[int]		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		str=13		,				UpperCAmelCase__    :		List[Any]=7		,				UpperCAmelCase__    :		List[Any]=True		,				UpperCAmelCase__    :		List[Any]=True		,				UpperCAmelCase__    :		Optional[int]=True		,				UpperCAmelCase__    :		Union[str, Any]=True		,				UpperCAmelCase__    :		Any=99		,				UpperCAmelCase__    :		Union[str, Any]=32		,				UpperCAmelCase__    :		List[str]=5		,				UpperCAmelCase__    :		Any=4		,				UpperCAmelCase__    :		List[Any]=37		,				UpperCAmelCase__    :		Optional[Any]="gelu"		,				UpperCAmelCase__    :		List[Any]=0.1		,				UpperCAmelCase__    :		Dict=0.1		,				UpperCAmelCase__    :		List[Any]=512		,				UpperCAmelCase__    :		Optional[Any]=16		,				UpperCAmelCase__    :		int=2		,				UpperCAmelCase__    :		Dict=0.02		,				UpperCAmelCase__    :		Any=3		,				UpperCAmelCase__    :		List[Any]=4		,				UpperCAmelCase__    :		Optional[Any]=None		,				):
										'''simple docstring'''
										lowercase   :       Dict								=parent
										lowercase   :       str								=batch_size
										lowercase   :       Optional[Any]								=seq_length
										lowercase   :       List[str]								=is_training
										lowercase   :       List[Any]								=use_input_mask
										lowercase   :       Union[str, Any]								=use_token_type_ids
										lowercase   :       Optional[Any]								=use_labels
										lowercase   :       Tuple								=vocab_size
										lowercase   :       Union[str, Any]								=hidden_size
										lowercase   :       List[Any]								=num_hidden_layers
										lowercase   :       str								=num_attention_heads
										lowercase   :       int								=intermediate_size
										lowercase   :       Union[str, Any]								=hidden_act
										lowercase   :       Dict								=hidden_dropout_prob
										lowercase   :       List[Any]								=attention_probs_dropout_prob
										lowercase   :       Union[str, Any]								=max_position_embeddings
										lowercase   :       Optional[int]								=type_vocab_size
										lowercase   :       Dict								=type_sequence_label_size
										lowercase   :       Dict								=initializer_range
										lowercase   :       str								=num_labels
										lowercase   :       int								=num_choices
										lowercase   :       Dict								=scope







				def    lowerCamelCase_  (					self    :		Tuple      ):
										'''simple docstring'''
										lowercase   :       Tuple								=ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

										lowercase   :       int								=None
										if self.use_token_type_ids:
																lowercase   :       Union[str, Any]								=ids_tensor([self.batch_size, self.seq_length]		,				self.type_vocab_size      )

										lowercase   :       Dict								=None
										lowercase   :       List[Any]								=None
										lowercase   :       List[str]								=None
										if self.use_labels:
																lowercase   :       Any								=ids_tensor([self.batch_size]		,				self.type_sequence_label_size      )
																lowercase   :       str								=ids_tensor([self.batch_size, self.seq_length]		,				self.num_labels      )
																lowercase   :       Tuple								=ids_tensor([self.batch_size]		,				self.num_choices      )

										lowercase   :       Optional[int]								=self.get_config()

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







				def    lowerCamelCase_  (					self    :		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=UpperCAmelCase__		,				initializer_range=self.initializer_range		,				)







				def    lowerCamelCase_  (					self    :		Optional[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Optional[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=NystromformerModel(config=UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       int								=model(UpperCAmelCase__		,				attention_mask=UpperCAmelCase__		,				token_type_ids=UpperCAmelCase__      )
										lowercase   :       Any								=model(UpperCAmelCase__		,				token_type_ids=UpperCAmelCase__      )
										lowercase   :       Any								=model(UpperCAmelCase__      )
										self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)      )







				def    lowerCamelCase_  (					self    :		str		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		Optional[Any]		,				UpperCAmelCase__    :		List[str]		,				UpperCAmelCase__    :		Any      ):
										'''simple docstring'''
										lowercase   :       Tuple								=NystromformerForMaskedLM(config=UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       Tuple								=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    :		List[Any]		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		List[Any]      ):
										'''simple docstring'''
										lowercase   :       Optional[Any]								=NystromformerForQuestionAnswering(config=UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       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    lowerCamelCase_  (					self    :		str		,				UpperCAmelCase__    :		Any		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		List[Any]      ):
										'''simple docstring'''
										lowercase   :       Optional[Any]								=self.num_labels
										lowercase   :       Any								=NystromformerForSequenceClassification(UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       Optional[int]								=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    :		List[Any]		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		List[str]		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       List[str]								=self.num_labels
										lowercase   :       Optional[Any]								=NystromformerForTokenClassification(config=UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       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.num_labels)      )







				def    lowerCamelCase_  (					self    :		Tuple		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Any		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		List[Any]      ):
										'''simple docstring'''
										lowercase   :       Any								=self.num_choices
										lowercase   :       List[str]								=NystromformerForMultipleChoice(config=UpperCAmelCase__      )
										model.to(UpperCAmelCase__      )
										model.eval()
										lowercase   :       List[Any]								=input_ids.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
										lowercase   :       int								=token_type_ids.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
										lowercase   :       Optional[int]								=input_mask.unsqueeze(1      ).expand(-1		,				self.num_choices		,				-1      ).contiguous()
										lowercase   :       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    lowerCamelCase_  (					self    :		Any      ):
										'''simple docstring'''
										lowercase   :       Dict								=self.prepare_config_and_inputs()
										(
										    (
										    lowercase
										)		,  (
										    lowercase
										)		,  (
										    lowercase
										)		,  (
										    lowercase
										)		,  (
										    lowercase
										)		,  (
										    lowercase
										)		,  (
										    lowercase
										)		,  
										)   :       Union[str, Any]								=config_and_inputs
										lowercase   :       Dict								={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
										return config, inputs_dict





@require_torch
class 		__SCREAMING_SNAKE_CASE       ( lowercase__	,       lowercase__	,       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    lowerCamelCase_  (					self    :		int      ):
										'''simple docstring'''
										lowercase   :       List[Any]								=NystromformerModelTester(self      )
										lowercase   :       Tuple								=ConfigTester(self		,				config_class=UpperCAmelCase__		,				hidden_size=37      )







				def    lowerCamelCase_  (					self    :		Tuple      ):
										'''simple docstring'''
										self.config_tester.run_common_tests()







				def    lowerCamelCase_  (					self    :		List[str]      ):
										'''simple docstring'''
										lowercase   :       Tuple								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		int      ):
										'''simple docstring'''
										lowercase   :       Dict								=self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																lowercase   :       Dict								=type
																self.model_tester.create_and_check_model(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Optional[int]      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										lowercase   :       Any								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										lowercase   :       Tuple								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Tuple      ):
										'''simple docstring'''
										lowercase   :       int								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Any      ):
										'''simple docstring'''
										lowercase   :       Union[str, Any]								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__      )







				@slow
				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																lowercase   :       List[Any]								=NystromformerModel.from_pretrained(UpperCAmelCase__      )
																self.assertIsNotNone(UpperCAmelCase__      )





@require_torch
class 		__SCREAMING_SNAKE_CASE       ( unittest.TestCase			):







				@slow
				def    lowerCamelCase_  (					self    :		Union[str, Any]      ):
										'''simple docstring'''
										lowercase   :       Dict								=NystromformerModel.from_pretrained('''uw-madison/nystromformer-512'''      )
										lowercase   :       Any								=torch.tensor([[0, 1, 2, 3, 4, 5]]      )

										with torch.no_grad():
																lowercase   :       Tuple								=model(UpperCAmelCase__      )[0]

										lowercase   :       int								=torch.Size((1, 6, 768)      )
										self.assertEqual(output.shape		,				UpperCAmelCase__      )

										lowercase   :       Dict								=torch.tensor(
										    [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]      )

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







				@slow
				def    lowerCamelCase_  (					self    :		Any      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								='''the [MASK] of Belgium is Brussels'''

										lowercase   :       Dict								=AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512'''      )
										lowercase   :       Optional[int]								=NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512'''      )

										lowercase   :       List[Any]								=tokenizer(UpperCAmelCase__		,				return_tensors='''pt'''      )

										with torch.no_grad():
																lowercase   :       Optional[int]								=model(encoding.input_ids      ).logits

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

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








'''simple docstring'''





import unittest

import numpy as np

from transformers import RobertaConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
							from transformers.models.roberta.modeling_flax_roberta import (
							    FlaxRobertaForCausalLM,
							    FlaxRobertaForMaskedLM,
							    FlaxRobertaForMultipleChoice,
							    FlaxRobertaForQuestionAnswering,
							    FlaxRobertaForSequenceClassification,
							    FlaxRobertaForTokenClassification,
							    FlaxRobertaModel,
							)





class 		__SCREAMING_SNAKE_CASE       ( unittest.TestCase			):







				def __init__(					self    :		Optional[Any]		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		str=13		,				UpperCAmelCase__    :		Optional[Any]=7		,				UpperCAmelCase__    :		Dict=True		,				UpperCAmelCase__    :		int=True		,				UpperCAmelCase__    :		int=True		,				UpperCAmelCase__    :		Dict=True		,				UpperCAmelCase__    :		Dict=99		,				UpperCAmelCase__    :		str=32		,				UpperCAmelCase__    :		Optional[Any]=5		,				UpperCAmelCase__    :		Any=4		,				UpperCAmelCase__    :		List[Any]=37		,				UpperCAmelCase__    :		str="gelu"		,				UpperCAmelCase__    :		Dict=0.1		,				UpperCAmelCase__    :		List[Any]=0.1		,				UpperCAmelCase__    :		Optional[Any]=512		,				UpperCAmelCase__    :		Any=16		,				UpperCAmelCase__    :		str=2		,				UpperCAmelCase__    :		Tuple=0.02		,				UpperCAmelCase__    :		int=4		,				):
										'''simple docstring'''
										lowercase   :       int								=parent
										lowercase   :       List[str]								=batch_size
										lowercase   :       str								=seq_length
										lowercase   :       Optional[Any]								=is_training
										lowercase   :       Union[str, Any]								=use_attention_mask
										lowercase   :       Optional[Any]								=use_token_type_ids
										lowercase   :       Tuple								=use_labels
										lowercase   :       List[str]								=vocab_size
										lowercase   :       List[str]								=hidden_size
										lowercase   :       Tuple								=num_hidden_layers
										lowercase   :       Any								=num_attention_heads
										lowercase   :       List[str]								=intermediate_size
										lowercase   :       Optional[Any]								=hidden_act
										lowercase   :       Dict								=hidden_dropout_prob
										lowercase   :       List[Any]								=attention_probs_dropout_prob
										lowercase   :       Optional[Any]								=max_position_embeddings
										lowercase   :       Tuple								=type_vocab_size
										lowercase   :       Optional[int]								=type_sequence_label_size
										lowercase   :       Optional[Any]								=initializer_range
										lowercase   :       Optional[int]								=num_choices







				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										lowercase   :       Optional[Any]								=ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

										lowercase   :       Union[str, Any]								=None
										if self.use_attention_mask:
																lowercase   :       List[str]								=random_attention_mask([self.batch_size, self.seq_length]      )

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

										lowercase   :       int								=RobertaConfig(
										    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		,				)

										return config, input_ids, token_type_ids, attention_mask







				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										lowercase   :       List[Any]								=self.prepare_config_and_inputs()
										lowercase		,  lowercase		,  lowercase		,  lowercase   :       str								=config_and_inputs
										lowercase   :       Optional[Any]								={'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
										return config, inputs_dict







				def    lowerCamelCase_  (					self    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       List[str]								=self.prepare_config_and_inputs()
										lowercase		,  lowercase		,  lowercase		,  lowercase   :       Any								=config_and_inputs

										lowercase   :       List[str]								=True
										lowercase   :       Optional[Any]								=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]      )
										lowercase   :       str								=ids_tensor([self.batch_size, self.seq_length]		,				vocab_size=2      )

										return (
										    config,
										    input_ids,
										    token_type_ids,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)





@require_flax
class 		__SCREAMING_SNAKE_CASE       ( lowercase__	,       unittest.TestCase			):
				lowerCamelCase_							=						True

				lowerCamelCase_							=						(
				    (
				        FlaxRobertaModel,
				        FlaxRobertaForCausalLM,
				        FlaxRobertaForMaskedLM,
				        FlaxRobertaForSequenceClassification,
				        FlaxRobertaForTokenClassification,
				        FlaxRobertaForMultipleChoice,
				        FlaxRobertaForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)







				def    lowerCamelCase_  (					self    :		List[Any]      ):
										'''simple docstring'''
										lowercase   :       str								=FlaxRobertaModelTester(self      )







				@slow
				def    lowerCamelCase_  (					self    :		Optional[int]      ):
										'''simple docstring'''
										for model_class_name in self.all_model_classes:
																lowercase   :       Optional[int]								=model_class_name.from_pretrained('''roberta-base'''		,				from_pt=UpperCAmelCase__      )
																lowercase   :       List[Any]								=model(np.ones((1, 1)      )      )
																self.assertIsNotNone(UpperCAmelCase__      )





