






import unittest

from transformers import MraConfig, 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, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
						import torch

						from transformers import (
						    MraForMaskedLM,
						    MraForMultipleChoice,
						    MraForQuestionAnswering,
						    MraForSequenceClassification,
						    MraForTokenClassification,
						    MraModel,
						)
						from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST


class        __UpperCamelCase    :







	def __init__(  self      :							Any       ,				_lowerCAmelCase      :							List[str]       ,				_lowerCAmelCase      :							Any=2       ,				_lowerCAmelCase      :							Tuple=8       ,				_lowerCAmelCase      :							List[Any]=True       ,				_lowerCAmelCase      :							Union[str, Any]=True       ,				_lowerCAmelCase      :							List[str]=True       ,				_lowerCAmelCase      :							Optional[int]=True       ,				_lowerCAmelCase      :							List[Any]=99       ,				_lowerCAmelCase      :							Tuple=16       ,				_lowerCAmelCase      :							Optional[Any]=5       ,				_lowerCAmelCase      :							Optional[Any]=2       ,				_lowerCAmelCase      :							Optional[Any]=36       ,				_lowerCAmelCase      :							Dict="gelu"       ,				_lowerCAmelCase      :							List[str]=0.0       ,				_lowerCAmelCase      :							List[Any]=0.0       ,				_lowerCAmelCase      :							Union[str, Any]=512       ,				_lowerCAmelCase      :							Optional[Any]=16       ,				_lowerCAmelCase      :							Optional[Any]=2       ,				_lowerCAmelCase      :							List[str]=0.02       ,				_lowerCAmelCase      :							str=3       ,				_lowerCAmelCase      :							Dict=4       ,				_lowerCAmelCase      :							Optional[int]=None       ,				)		-> Dict:
					"""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			)		-> Optional[int]:
					"""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      :							Any			)		-> int:
					"""simple docstring"""







					return MraConfig(
					    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=_lowerCAmelCase       ,				initializer_range=self.initializer_range       ,				)







	def  _a   (  self      :							Union[str, Any]			)		-> List[str]:
					"""simple docstring"""







					__lowercase             =  self.get_config()
					__lowercase             =  300
					return config







	def  _a   (  self      :							List[str]			)		-> Dict:
					"""simple docstring"""







					(
					    (
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,						(
					    __lowercase						
					)       ,												
					)             =  self.prepare_config_and_inputs()

					__lowercase             =  True
					__lowercase             =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
					__lowercase             =  ids_tensor([self.batch_size, self.seq_length]       ,				vocab_size=2			)

					return (
					    config,
					    input_ids,
					    token_type_ids,
					    input_mask,
					    sequence_labels,
					    token_labels,
					    choice_labels,
					    encoder_hidden_states,
					    encoder_attention_mask,
					)







	def  _a   (  self      :							int       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							Optional[int]       ,				_lowerCAmelCase      :							Union[str, Any]       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							Any			)		-> List[Any]:
					"""simple docstring"""







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







	def  _a   (  self      :							int       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							Union[str, Any]       ,				_lowerCAmelCase      :							Optional[Any]       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							Union[str, Any]       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							List[str]       ,				)		-> str:
					"""simple docstring"""







					__lowercase             =  True
					__lowercase             =  MraModel(_lowerCAmelCase			)
					model.to(_lowerCAmelCase			)
					model.eval()
					__lowercase             =  model(
					    _lowerCAmelCase       ,				attention_mask=_lowerCAmelCase       ,				token_type_ids=_lowerCAmelCase       ,				encoder_hidden_states=_lowerCAmelCase       ,				encoder_attention_mask=_lowerCAmelCase       ,				)
					__lowercase             =  model(
					    _lowerCAmelCase       ,				attention_mask=_lowerCAmelCase       ,				token_type_ids=_lowerCAmelCase       ,				encoder_hidden_states=_lowerCAmelCase       ,				)
					__lowercase             =  model(_lowerCAmelCase       ,				attention_mask=_lowerCAmelCase       ,				token_type_ids=_lowerCAmelCase			)
					self.parent.assertEqual(result.last_hidden_state.shape       ,				(self.batch_size, self.seq_length, self.hidden_size)			)







	def  _a   (  self      :							Tuple       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							List[Any]			)		-> Dict:
					"""simple docstring"""







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







	def  _a   (  self      :							Any       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							int			)		-> Tuple:
					"""simple docstring"""







					__lowercase             =  MraForQuestionAnswering(config=_lowerCAmelCase			)
					model.to(_lowerCAmelCase			)
					model.eval()
					__lowercase             =  model(
					    _lowerCAmelCase       ,				attention_mask=_lowerCAmelCase       ,				token_type_ids=_lowerCAmelCase       ,				start_positions=_lowerCAmelCase       ,				end_positions=_lowerCAmelCase       ,				)
					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      :							Union[str, Any]       ,				_lowerCAmelCase      :							Optional[int]       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							List[str]       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							Optional[int]			)		-> Union[str, Any]:
					"""simple docstring"""







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







	def  _a   (  self      :							Any       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							str       ,				_lowerCAmelCase      :							Optional[Any]       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							List[Any]			)		-> Union[str, Any]:
					"""simple docstring"""







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







	def  _a   (  self      :							Any       ,				_lowerCAmelCase      :							List[Any]       ,				_lowerCAmelCase      :							Dict       ,				_lowerCAmelCase      :							Union[str, Any]       ,				_lowerCAmelCase      :							int       ,				_lowerCAmelCase      :							Any       ,				_lowerCAmelCase      :							Tuple       ,				_lowerCAmelCase      :							Tuple			)		-> List[str]:
					"""simple docstring"""







					__lowercase             =  self.num_choices
					__lowercase             =  MraForMultipleChoice(config=_lowerCAmelCase			)
					model.to(_lowerCAmelCase			)
					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(
					    _lowerCAmelCase       ,				attention_mask=_lowerCAmelCase       ,				token_type_ids=_lowerCAmelCase       ,				labels=_lowerCAmelCase       ,				)
					self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.num_choices)			)







	def  _a   (  self      :							Optional[int]			)		-> 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        __UpperCamelCase    (     _lowerCAmelCase   ,			unittest.TestCase   ):
	__snake_case	:Dict   =       (
	    (
	        MraModel,
	        MraForMaskedLM,
	        MraForMultipleChoice,
	        MraForQuestionAnswering,
	        MraForSequenceClassification,
	        MraForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	__snake_case	:List[Any]   =       False
	__snake_case	:List[str]   =       False
	__snake_case	:Dict   =       False
	__snake_case	:Tuple   =       False

	__snake_case	:Optional[Any]   =       ()







	def  _a   (  self      :							Tuple			)		-> Optional[int]:
					"""simple docstring"""







					__lowercase             =  MraModelTester(self			)
					__lowercase             =  ConfigTester(self       ,				config_class=_lowerCAmelCase       ,				hidden_size=37			)







	def  _a   (  self      :							Any			)		-> Any:
					"""simple docstring"""







					self.config_tester.run_common_tests()







	def  _a   (  self      :							Any			)		-> Any:
					"""simple docstring"""







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







	def  _a   (  self      :							Dict			)		-> 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(*_lowerCAmelCase			)







	def  _a   (  self      :							int			)		-> List[str]:
					"""simple docstring"""







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







	def  _a   (  self      :							List[str]			)		-> List[Any]:
					"""simple docstring"""







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







	def  _a   (  self      :							int			)		-> int:
					"""simple docstring"""







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







	def  _a   (  self      :							List[Any]			)		-> Tuple:
					"""simple docstring"""







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







	def  _a   (  self      :							int			)		-> Dict:
					"""simple docstring"""







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







	@slow
	def  _a   (  self      :							List[Any]			)		-> Union[str, Any]:
					"""simple docstring"""







					for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowercase             =  MraModel.from_pretrained(_lowerCAmelCase			)
									self.assertIsNotNone(_lowerCAmelCase			)







	@unittest.skip(reason="""MRA does not output attentions"""			)
	def  _a   (  self      :							List[str]			)		-> List[Any]:
					"""simple docstring"""







					return


@require_torch
class        __UpperCamelCase    (     unittest.TestCase   ):







	@slow
	def  _a   (  self      :							List[str]			)		-> List[Any]:
					"""simple docstring"""







					__lowercase             =  MraModel.from_pretrained("""uw-madison/mra-base-512-4"""			)
					__lowercase             =  torch.arange(256			).unsqueeze(0			)

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

					__lowercase             =  torch.Size((1, 256, 768)			)
					self.assertEqual(output.shape       ,				_lowerCAmelCase			)

					__lowercase             =  torch.tensor(
					    [[[-0.0_140, 0.0_830, -0.0_381], [0.1_546, 0.1_402, 0.0_220], [0.1_162, 0.0_851, 0.0_165]]]			)

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







	@slow
	def  _a   (  self      :							str			)		-> Union[str, Any]:
					"""simple docstring"""







					__lowercase             =  MraForMaskedLM.from_pretrained("""uw-madison/mra-base-512-4"""			)
					__lowercase             =  torch.arange(256			).unsqueeze(0			)

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

					__lowercase             =  5_0265

					__lowercase             =  torch.Size((1, 256, vocab_size)			)
					self.assertEqual(output.shape       ,				_lowerCAmelCase			)

					__lowercase             =  torch.tensor(
					    [[[9.2_595, -3.6_038, 11.8_819], [9.3_869, -3.2_693, 11.0_956], [11.8_524, -3.4_938, 13.1_210]]]			)

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







	@slow
	def  _a   (  self      :							Optional[int]			)		-> List[str]:
					"""simple docstring"""







					__lowercase             =  MraForMaskedLM.from_pretrained("""uw-madison/mra-base-4096-8-d3"""			)
					__lowercase             =  torch.arange(4096			).unsqueeze(0			)

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

					__lowercase             =  5_0265

					__lowercase             =  torch.Size((1, 4096, vocab_size)			)
					self.assertEqual(output.shape       ,				_lowerCAmelCase			)

					__lowercase             =  torch.tensor(
					    [[[5.4_789, -2.3_564, 7.5_064], [7.9_067, -1.3_369, 9.9_668], [9.0_712, -1.8_106, 7.0_380]]]			)

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









from math import sqrt





def 				snake_case  (     lowerCamelCase					):




	'''simple docstring'''







	if 1 < number < 4:
		# 2 and 3 are primes
		return True
	elif number < 2 or number % 2 == 0 or number % 3 == 0:
		# Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
		return False

	# All primes number are in format of 6k +/- 1
	for i in range(5				,  int(sqrt(lowerCamelCase					) + 1					)				,  6					):
		if number % i == 0 or number % (i + 2) == 0:
			return False
	return True





def 				snake_case  (     lowerCamelCase = 10_001					):




	'''simple docstring'''







	__lowercase             =  0
	__lowercase             =  1
	while count != nth and number < 3:
		number += 1
		if is_prime(lowerCamelCase					):
			count += 1
	while count != nth:
		number += 2
		if is_prime(lowerCamelCase					):
			count += 1
	return number


if __name__ == "__main__":
						print(F'''{solution() = }''')

