



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







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







					def  _lowercase			(       self      :     Optional[Any]):
											A__	:							List[Any]						    = ids_tensor([self.batch_size, self.seq_length]			,		self.vocab_size)

											A__	:							Optional[int]						    = None
											if self.use_input_mask:
																	A__	:							Union[str, Any]						    = random_attention_mask([self.batch_size, self.seq_length])

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

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

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

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







					def  _lowercase			(       self      :     Any):
											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=_A			,		initializer_range=self.initializer_range			,		)







					def  _lowercase			(       self      :     Tuple			,		_A      :     int			,		_A      :     int			,		_A      :     Tuple			,		_A      :     Any			,		_A      :     Union[str, Any]			,		_A      :     str			,		_A      :     Dict):
											A__	:							Dict						    = NystromformerModel(config=_A)
											model.to(_A)
											model.eval()
											A__	:							Optional[Any]						    = model(_A			,		attention_mask=_A			,		token_type_ids=_A)
											A__	:							Optional[int]						    = model(_A			,		token_type_ids=_A)
											A__	:							str						    = model(_A)
											self.parent.assertEqual(result.last_hidden_state.shape			,		(self.batch_size, self.seq_length, self.hidden_size))







					def  _lowercase			(       self      :     Optional[Any]			,		_A      :     Optional[int]			,		_A      :     Any			,		_A      :     List[Any]			,		_A      :     Optional[int]			,		_A      :     Optional[Any]			,		_A      :     List[str]			,		_A      :     Dict):
											A__	:							Union[str, Any]						    = NystromformerForMaskedLM(config=_A)
											model.to(_A)
											model.eval()
											A__	:							Dict						    = model(_A			,		attention_mask=_A			,		token_type_ids=_A			,		labels=_A)
											self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.seq_length, self.vocab_size))







					def  _lowercase			(       self      :     List[str]			,		_A      :     Union[str, Any]			,		_A      :     str			,		_A      :     Tuple			,		_A      :     List[Any]			,		_A      :     Tuple			,		_A      :     Optional[int]			,		_A      :     Optional[int]):
											A__	:							Union[str, Any]						    = NystromformerForQuestionAnswering(config=_A)
											model.to(_A)
											model.eval()
											A__	:							Tuple						    = model(
											    _A			,		attention_mask=_A			,		token_type_ids=_A			,		start_positions=_A			,		end_positions=_A			,		)
											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  _lowercase			(       self      :     Optional[int]			,		_A      :     Optional[Any]			,		_A      :     Tuple			,		_A      :     List[str]			,		_A      :     Tuple			,		_A      :     Any			,		_A      :     Optional[Any]			,		_A      :     Tuple):
											A__	:							Optional[Any]						    = self.num_labels
											A__	:							Any						    = NystromformerForSequenceClassification(_A)
											model.to(_A)
											model.eval()
											A__	:							Union[str, Any]						    = model(_A			,		attention_mask=_A			,		token_type_ids=_A			,		labels=_A)
											self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.num_labels))







					def  _lowercase			(       self      :     Dict			,		_A      :     int			,		_A      :     int			,		_A      :     Tuple			,		_A      :     Dict			,		_A      :     Optional[Any]			,		_A      :     List[Any]			,		_A      :     Dict):
											A__	:							Any						    = self.num_labels
											A__	:							str						    = NystromformerForTokenClassification(config=_A)
											model.to(_A)
											model.eval()
											A__	:							Any						    = model(_A			,		attention_mask=_A			,		token_type_ids=_A			,		labels=_A)
											self.parent.assertEqual(result.logits.shape			,		(self.batch_size, self.seq_length, self.num_labels))







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







					def  _lowercase			(       self      :     Union[str, Any]):
											A__	:							Any						    = self.prepare_config_and_inputs()
											(
											    (
											    A__
											)    ,       (
											    A__
											)    ,       (
											    A__
											)    ,       (
											    A__
											)    ,       (
											    A__
											)    ,       (
											    A__
											)    ,       (
											    A__
											)    ,       
											)	:							Tuple						    = config_and_inputs
											A__	:							Any						    = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
											return config, inputs_dict




@require_torch
class 							lowerCAmelCase__  (	UpperCamelCase					,							UpperCamelCase					,							unittest.TestCase						):
					__A     :    Optional[Any]           =						(
					    (
					        NystromformerModel,
					        NystromformerForMaskedLM,
					        NystromformerForMultipleChoice,
					        NystromformerForQuestionAnswering,
					        NystromformerForSequenceClassification,
					        NystromformerForTokenClassification,
					    )
					    if is_torch_available()
					    else ()
					)
					__A     :    str           =						(
					    {
					        'feature-extraction': NystromformerModel,
					        'fill-mask': NystromformerForMaskedLM,
					        'question-answering': NystromformerForQuestionAnswering,
					        'text-classification': NystromformerForSequenceClassification,
					        'token-classification': NystromformerForTokenClassification,
					        'zero-shot': NystromformerForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					__A     :    Union[str, Any]           =						False
					__A     :    Optional[Any]           =						False







					def  _lowercase			(       self      :     Tuple):
											A__	:							Optional[Any]						    = NystromformerModelTester(self)
											A__	:							Any						    = ConfigTester(self			,		config_class=_A			,		hidden_size=37)







					def  _lowercase			(       self      :     List[Any]):
											self.config_tester.run_common_tests()







					def  _lowercase			(       self      :     int):
											A__	:							Tuple						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_A)







					def  _lowercase			(       self      :     Union[str, Any]):
											A__	:							List[Any]						    = self.model_tester.prepare_config_and_inputs()
											for type in ["absolute", "relative_key", "relative_key_query"]:
																	A__	:							str						    = type
																	self.model_tester.create_and_check_model(*_A)







					def  _lowercase			(       self      :     int):
											A__	:							Optional[Any]						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*_A)







					def  _lowercase			(       self      :     str):
											A__	:							Dict						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_multiple_choice(*_A)







					def  _lowercase			(       self      :     List[Any]):
											A__	:							Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_question_answering(*_A)







					def  _lowercase			(       self      :     List[Any]):
											A__	:							str						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_sequence_classification(*_A)







					def  _lowercase			(       self      :     Dict):
											A__	:							Union[str, Any]						    = self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*_A)







					@slow
					def  _lowercase			(       self      :     Optional[int]):
											for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	A__	:							str						    = NystromformerModel.from_pretrained(_A)
																	self.assertIsNotNone(_A)




@require_torch
class 							lowerCAmelCase__  (	unittest.TestCase						):







					@slow
					def  _lowercase			(       self      :     Optional[int]):
											A__	:							Union[str, Any]						    = NystromformerModel.from_pretrained("uw-madison/nystromformer-512")
											A__	:							Union[str, Any]						    = torch.tensor([[0, 1, 2, 3, 4, 5]])

											with torch.no_grad():
																	A__	:							Union[str, Any]						    = model(_A)[0]

											A__	:							Optional[Any]						    = torch.Size((1, 6, 768))
											self.assertEqual(output.shape			,		_A)

											A__	:							Tuple						    = 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]			,		_A			,		atol=1e-4))







					@slow
					def  _lowercase			(       self      :     Dict):
											A__	:							Dict						    = "the [MASK] of Belgium is Brussels"

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

											A__	:							Optional[int]						    = tokenizer(_A			,		return_tensors="pt")

											with torch.no_grad():
																	A__	:							Union[str, Any]						    = model(encoding.input_ids).logits

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

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




snake_case    :						Tuple  =						'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
def 	snake_case__       (					)		->		None:







						"""simple docstring"""

						A__	:							Union[str, Any]						    = input("Enter message: "       )
						A__	:							Tuple						    = input("Enter key [alphanumeric]: "       )
						A__	:							Optional[Any]						    = input("Encrypt/Decrypt [e/d]: "       )

						if mode.lower().startswith("e"       ):
												A__	:							Tuple						    = "encrypt"
												A__	:							str						    = encrypt_message(__lowercase							,     __lowercase       )
						elif mode.lower().startswith("d"       ):
												A__	:							Optional[Any]						    = "decrypt"
												A__	:							Union[str, Any]						    = decrypt_message(__lowercase							,     __lowercase       )

						print(F'\n{mode.title()}ed message:'       )
						print(__lowercase       )
def 	snake_case__       (					__lowercase							,     __lowercase       )		->		str:







						"""simple docstring"""

						return translate_message(__lowercase							,     __lowercase							,     "encrypt"       )
def 	snake_case__       (					__lowercase							,     __lowercase       )		->		str:







						"""simple docstring"""

						return translate_message(__lowercase							,     __lowercase							,     "decrypt"       )
def 	snake_case__       (					__lowercase							,     __lowercase							,     __lowercase       )		->		str:







						"""simple docstring"""

						A__	:							Dict						    = []
						A__	:							Union[str, Any]						    = 0
						A__	:							List[Any]						    = key.upper()

						for symbol in message:
												A__	:							List[str]						    = LETTERS.find(symbol.upper()       )
												if num != -1:
																		if mode == "encrypt":
																								num += LETTERS.find(key[key_index]       )
																		elif mode == "decrypt":
																								num -= LETTERS.find(key[key_index]       )

																		num %= len(__lowercase       )

																		if symbol.isupper():
																								translated.append(LETTERS[num]       )
																		elif symbol.islower():
																								translated.append(LETTERS[num].lower()       )

																		key_index += 1
																		if key_index == len(__lowercase       ):
																								A__	:							Optional[int]						    = 0
												else:
																		translated.append(__lowercase       )
						return "".join(__lowercase       )


if __name__ == "__main__":
		main()