






"""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    __lowerCamelCase    :

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

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[int]:
       UpperCamelCase__         =    ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size      )

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

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

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

       UpperCamelCase__         =    self.get_config()

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

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       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 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       str:
       UpperCamelCase__         =    NystromformerModel(config=A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    model(A__ ,				attention_mask=A__ ,				token_type_ids=A__      )
       UpperCamelCase__         =    model(A__ ,				token_type_ids=A__      )
       UpperCamelCase__         =    model(A__      )
       self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)      )

   def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       Optional[int]:
       UpperCamelCase__         =    NystromformerForMaskedLM(config=A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       List[Any]:
       UpperCamelCase__         =    NystromformerForQuestionAnswering(config=A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       Dict:
       UpperCamelCase__         =    self.num_labels
       UpperCamelCase__         =    NystromformerForSequenceClassification(A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    model(A__ ,				attention_mask=A__ ,				token_type_ids=A__ ,				labels=A__      )
       self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)      )

   def 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       Any:
       UpperCamelCase__         =    self.num_labels
       UpperCamelCase__         =    NystromformerForTokenClassification(config=A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    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 						SCREAMING_SNAKE_CASE__ (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_      )		->       int:
       UpperCamelCase__         =    self.num_choices
       UpperCamelCase__         =    NystromformerForMultipleChoice(config=A__      )
       model.to(A__      )
       model.eval()
       UpperCamelCase__         =    input_ids.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
       UpperCamelCase__         =    token_type_ids.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
       UpperCamelCase__         =    input_mask.unsqueeze(1      ).expand(-1 ,				self.num_choices ,				-1      ).contiguous()
       UpperCamelCase__         =    model(
           A__ ,				attention_mask=A__ ,				token_type_ids=A__ ,				labels=A__ ,				)
       self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)      )







   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       List[Any]:
       UpperCamelCase__         =    self.prepare_config_and_inputs()
       (
           (
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,						(
           UpperCamelCase__		
       )       ,								
       )         =    config_and_inputs
       UpperCamelCase__         =    {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
       return config, inputs_dict







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

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    NystromformerModelTester(self      )
       UpperCamelCase__         =    ConfigTester(self ,				config_class=A__ ,				hidden_size=37      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Any:
       self.config_tester.run_common_tests()

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       str:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       for type in ["absolute", "relative_key", "relative_key_query"]:
           UpperCamelCase__         =    type
           self.model_tester.create_and_check_model(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       List[Any]:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[Any]:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_multiple_choice(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Tuple:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_question_answering(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Optional[Any]:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_sequence_classification(*A__      )

   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Dict:
       UpperCamelCase__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A__      )







   @slow
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Tuple:
       for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           UpperCamelCase__         =    NystromformerModel.from_pretrained(A__      )
           self.assertIsNotNone(A__      )







@require_torch
class    __lowerCamelCase    (  unittest.TestCase  ):

   @slow
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       List[str]:
       UpperCamelCase__         =    NystromformerModel.from_pretrained('uw-madison/nystromformer-512'      )
       UpperCamelCase__         =    torch.tensor([[0, 1, 2, 3, 4, 5]]      )

       with torch.no_grad():
           UpperCamelCase__         =    model(A__      )[0]

       UpperCamelCase__         =    torch.Size((1, 6, 768)      )
       self.assertEqual(output.shape ,				A__      )

       UpperCamelCase__         =    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] ,				A__ ,				atol=1E-4      )      )







   @slow
   def 						SCREAMING_SNAKE_CASE__ (					self      )		->       Union[str, Any]:
       UpperCamelCase__         =    'the [MASK] of Belgium is Brussels'

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

       UpperCamelCase__         =    tokenizer(A__ ,				return_tensors='pt'      )

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

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

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












"""simple docstring"""





def   lowerCAmelCase_( SCREAMING_SNAKE_CASE					)					->	bool:
    """simple docstring"""

    UpperCamelCase__         =    (1 + 24 * n) ** 0.5
    return ((1 + root) / 6) % 1 == 0







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE = 50_00					)					->	int:
    """simple docstring"""

    UpperCamelCase__         =    [(i * (3 * i - 1)) // 2 for i in range(1 ,	SCREAMING_SNAKE_CASE					)]
    for i, pentagonal_i in enumerate(SCREAMING_SNAKE_CASE					):
        for j in range(SCREAMING_SNAKE_CASE ,	len(SCREAMING_SNAKE_CASE					)					):
            UpperCamelCase__         =    pentagonal_nums[j]
            UpperCamelCase__         =    pentagonal_i + pentagonal_j
            UpperCamelCase__         =    pentagonal_j - pentagonal_i
            if is_pentagonal(SCREAMING_SNAKE_CASE					) and is_pentagonal(SCREAMING_SNAKE_CASE					):
                return b

    return -1


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




