




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







 '''simple docstring'''





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




   '''simple docstring'''



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


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						ids_tensor([self.batch_size, self.seq_length],   self.vocab_size	)

   _snake_case							:	Optional[int]			=						None
   if self.use_input_mask:
     _snake_case							:	Optional[int]			=						random_attention_mask([self.batch_size, self.seq_length]	)

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

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

   _snake_case							:	Any			=						self.get_config()

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


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


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	List[Any],   a_:	Any,   a_:	Dict,   a_:	Dict,   a_:	Any,   a_:	Dict,   a_:	Tuple	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						NystromformerModel(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	List[str]			=						model(a_,   attention_mask=a_,   token_type_ids=a_	)
   _snake_case							:	str			=						model(a_,   token_type_ids=a_	)
   _snake_case							:	int			=						model(a_	)
   self.parent.assertEqual(result.last_hidden_state.shape,   (self.batch_size, self.seq_length, self.hidden_size)	)


 def 				UpperCamelCase_     (       self:	Any,   a_:	int,   a_:	str,   a_:	Any,   a_:	List[Any],   a_:	Union[str, Any],   a_:	Union[str, Any],   a_:	str	):




   '''simple docstring'''



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




   '''simple docstring'''



   _snake_case							:	Tuple			=						NystromformerForQuestionAnswering(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Optional[int]			=						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 				UpperCamelCase_     (       self:	Union[str, Any],   a_:	Any,   a_:	Optional[int],   a_:	Tuple,   a_:	List[Any],   a_:	Dict,   a_:	Optional[int],   a_:	Any	):




   '''simple docstring'''



   _snake_case							:	int			=						self.num_labels
   _snake_case							:	str			=						NystromformerForSequenceClassification(a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Optional[int]			=						model(a_,   attention_mask=a_,   token_type_ids=a_,   labels=a_	)
   self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_labels)	)


 def 				UpperCamelCase_     (       self:	Optional[Any],   a_:	str,   a_:	Dict,   a_:	Dict,   a_:	List[Any],   a_:	List[Any],   a_:	str,   a_:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	int			=						self.num_labels
   _snake_case							:	List[Any]			=						NystromformerForTokenClassification(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	Tuple			=						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 				UpperCamelCase_     (       self:	Optional[Any],   a_:	Tuple,   a_:	str,   a_:	Optional[int],   a_:	Any,   a_:	Optional[Any],   a_:	int,   a_:	Dict	):




   '''simple docstring'''



   _snake_case							:	Any			=						self.num_choices
   _snake_case							:	Dict			=						NystromformerForMultipleChoice(config=a_	)
   model.to(a_	)
   model.eval()
   _snake_case							:	int			=						input_ids.unsqueeze(1	).expand(-1,   self.num_choices,   -1	).contiguous()
   _snake_case							:	Optional[int]			=						token_type_ids.unsqueeze(1	).expand(-1,   self.num_choices,   -1	).contiguous()
   _snake_case							:	int			=						input_mask.unsqueeze(1	).expand(-1,   self.num_choices,   -1	).contiguous()
   _snake_case							:	int			=						model(
       a_,   attention_mask=a_,   token_type_ids=a_,   labels=a_,   )
   self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_choices)	)


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self.prepare_config_and_inputs()
   (
       (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      (
       _snake_case
   )	,      
   )							:	List[Any]			=						config_and_inputs
   _snake_case							:	str			=						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
   return config, inputs_dict







@require_torch
class       lowercase(  __a ,    __a ,    unittest.TestCase				):







 '''simple docstring'''





 lowercase__				      =      (
     (
         NystromformerModel,
         NystromformerForMaskedLM,
         NystromformerForMultipleChoice,
         NystromformerForQuestionAnswering,
         NystromformerForSequenceClassification,
         NystromformerForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 lowercase__				      =      (
     {
         "feature-extraction": NystromformerModel,
         "fill-mask": NystromformerForMaskedLM,
         "question-answering": NystromformerForQuestionAnswering,
         "text-classification": NystromformerForSequenceClassification,
         "token-classification": NystromformerForTokenClassification,
         "zero-shot": NystromformerForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )
 lowercase__				      =      False
 lowercase__				      =      False


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   _snake_case							:	int			=						NystromformerModelTester(self	)
   _snake_case							:	str			=						ConfigTester(self,   config_class=a_,   hidden_size=37	)


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   self.config_tester.run_common_tests()


 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   _snake_case							:	List[str]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*a_	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	int			=						self.model_tester.prepare_config_and_inputs()
   for type in ["absolute", "relative_key", "relative_key_query"]:
     _snake_case							:	Any			=						type
     self.model_tester.create_and_check_model(*a_	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*a_	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	str			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_multiple_choice(*a_	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	int			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_question_answering(*a_	)


 def 				UpperCamelCase_     (       self:	Dict	):




   '''simple docstring'''



   _snake_case							:	str			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_sequence_classification(*a_	)


 def 				UpperCamelCase_     (       self:	str	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*a_	)


 @slow
 def 				UpperCamelCase_     (       self:	int	):




   '''simple docstring'''



   for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case							:	str			=						NystromformerModel.from_pretrained(a_	)
     self.assertIsNotNone(a_	)







@require_torch
class       lowercase(  unittest.TestCase				):







 '''simple docstring'''





 @slow
 def 				UpperCamelCase_     (       self:	List[Any]	):




   '''simple docstring'''



   _snake_case							:	Tuple			=						NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""	)
   _snake_case							:	List[str]			=						torch.tensor([[0, 1, 2, 3, 4, 5]]	)

   with torch.no_grad():
     _snake_case							:	Optional[Any]			=						model(a_	)[0]

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

   _snake_case							:	Optional[Any]			=						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 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						"""the [MASK] of Belgium is Brussels"""

   _snake_case							:	Optional[Any]			=						AutoTokenizer.from_pretrained("""uw-madison/nystromformer-512"""	)
   _snake_case							:	Optional[Any]			=						NystromformerForMaskedLM.from_pretrained("""uw-madison/nystromformer-512"""	)

   _snake_case							:	Optional[Any]			=						tokenizer(a_,   return_tensors="""pt"""	)

   with torch.no_grad():
     _snake_case							:	List[Any]			=						model(encoding.input_ids	).logits

   _snake_case							:	Optional[int]			=						token_logits[:, 2, :].argmax(-1	)[0]

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








"""simple docstring"""


import tempfile

import torch

from diffusers import IPNDMScheduler

from .test_schedulers import SchedulerCommonTest






class       lowercase(  __a				):







 '''simple docstring'''





 lowercase__				      =      (IPNDMScheduler,)
 lowercase__				      =      (("num_inference_steps", 50),)


 def 				UpperCamelCase_     (       self:	Union[str, Any],   **a_:	Union[str, Any]	):




   '''simple docstring'''



   _snake_case							:	List[Any]			=						{"""num_train_timesteps""": 1_000}
   config.update(**a_	)
   return config


 def 				UpperCamelCase_     (       self:	Tuple,   a_:	Optional[int]=0,   **a_:	int	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						dict(self.forward_default_kwargs	)
   _snake_case							:	Optional[Any]			=						kwargs.pop("""num_inference_steps""",   a_	)
   _snake_case							:	Optional[Any]			=						self.dummy_sample
   _snake_case							:	Dict			=						0.1 * sample
   _snake_case							:	Any			=						[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

   for scheduler_class in self.scheduler_classes:
     _snake_case							:	int			=						self.get_scheduler_config(**a_	)
     _snake_case							:	Dict			=						scheduler_class(**a_	)
     scheduler.set_timesteps(a_	)
     # copy over dummy past residuals
     _snake_case							:	int			=						dummy_past_residuals[:]

     if time_step is None:
       _snake_case							:	Union[str, Any]			=						scheduler.timesteps[len(scheduler.timesteps	) // 2]

     with tempfile.TemporaryDirectory() as tmpdirname:
       scheduler.save_config(a_	)
       _snake_case							:	Tuple			=						scheduler_class.from_pretrained(a_	)
       new_scheduler.set_timesteps(a_	)
       # copy over dummy past residuals
       _snake_case							:	Optional[Any]			=						dummy_past_residuals[:]

     _snake_case							:	List[Any]			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	str			=						new_scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     assert torch.sum(torch.abs(output - new_output	)	) < 1E-5, "Scheduler outputs are not identical"

     _snake_case							:	Optional[Any]			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	Optional[int]			=						new_scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     assert torch.sum(torch.abs(output - new_output	)	) < 1E-5, "Scheduler outputs are not identical"


 def 				UpperCamelCase_     (       self:	List[str]	):




   '''simple docstring'''



   pass


 def 				UpperCamelCase_     (       self:	str,   a_:	Any=0,   **a_:	Tuple	):




   '''simple docstring'''



   _snake_case							:	str			=						dict(self.forward_default_kwargs	)
   _snake_case							:	List[Any]			=						kwargs.pop("""num_inference_steps""",   a_	)
   _snake_case							:	Optional[int]			=						self.dummy_sample
   _snake_case							:	Tuple			=						0.1 * sample
   _snake_case							:	Optional[int]			=						[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

   for scheduler_class in self.scheduler_classes:
     _snake_case							:	Any			=						self.get_scheduler_config()
     _snake_case							:	Tuple			=						scheduler_class(**a_	)
     scheduler.set_timesteps(a_	)

     # copy over dummy past residuals (must be after setting timesteps)
     _snake_case							:	Union[str, Any]			=						dummy_past_residuals[:]

     if time_step is None:
       _snake_case							:	Tuple			=						scheduler.timesteps[len(scheduler.timesteps	) // 2]

     with tempfile.TemporaryDirectory() as tmpdirname:
       scheduler.save_config(a_	)
       _snake_case							:	List[str]			=						scheduler_class.from_pretrained(a_	)
       # copy over dummy past residuals
       new_scheduler.set_timesteps(a_	)

       # copy over dummy past residual (must be after setting timesteps)
       _snake_case							:	List[str]			=						dummy_past_residuals[:]

     _snake_case							:	str			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	Any			=						new_scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     assert torch.sum(torch.abs(output - new_output	)	) < 1E-5, "Scheduler outputs are not identical"

     _snake_case							:	str			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	int			=						new_scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     assert torch.sum(torch.abs(output - new_output	)	) < 1E-5, "Scheduler outputs are not identical"


 def 				UpperCamelCase_     (       self:	List[Any],   **a_:	Optional[int]	):




   '''simple docstring'''



   _snake_case							:	Union[str, Any]			=						self.scheduler_classes[0]
   _snake_case							:	Any			=						self.get_scheduler_config(**a_	)
   _snake_case							:	List[Any]			=						scheduler_class(**a_	)

   _snake_case							:	Union[str, Any]			=						10
   _snake_case							:	Union[str, Any]			=						self.dummy_model()
   _snake_case							:	List[Any]			=						self.dummy_sample_deter
   scheduler.set_timesteps(a_	)

   for i, t in enumerate(scheduler.timesteps	):
     _snake_case							:	Optional[Any]			=						model(a_,   a_	)
     _snake_case							:	Any			=						scheduler.step(a_,   a_,   a_	).prev_sample

   for i, t in enumerate(scheduler.timesteps	):
     _snake_case							:	Union[str, Any]			=						model(a_,   a_	)
     _snake_case							:	Any			=						scheduler.step(a_,   a_,   a_	).prev_sample

   return sample


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   _snake_case							:	str			=						dict(self.forward_default_kwargs	)

   _snake_case							:	int			=						kwargs.pop("""num_inference_steps""",   a_	)

   for scheduler_class in self.scheduler_classes:
     _snake_case							:	Union[str, Any]			=						self.get_scheduler_config()
     _snake_case							:	Tuple			=						scheduler_class(**a_	)

     _snake_case							:	Dict			=						self.dummy_sample
     _snake_case							:	List[str]			=						0.1 * sample

     if num_inference_steps is not None and hasattr(a_,   """set_timesteps"""	):
       scheduler.set_timesteps(a_	)
     elif num_inference_steps is not None and not hasattr(a_,   """set_timesteps"""	):
       _snake_case							:	Dict			=						num_inference_steps

     # copy over dummy past residuals (must be done after set_timesteps)
     _snake_case							:	Union[str, Any]			=						[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
     _snake_case							:	List[str]			=						dummy_past_residuals[:]

     _snake_case							:	Optional[int]			=						scheduler.timesteps[5]
     _snake_case							:	Optional[Any]			=						scheduler.timesteps[6]

     _snake_case							:	str			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	List[str]			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     self.assertEqual(output_a.shape,   sample.shape	)
     self.assertEqual(output_a.shape,   output_a.shape	)

     _snake_case							:	Any			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample
     _snake_case							:	Any			=						scheduler.step(a_,   a_,   a_,   **a_	).prev_sample

     self.assertEqual(output_a.shape,   sample.shape	)
     self.assertEqual(output_a.shape,   output_a.shape	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   for timesteps in [100, 1_000]:
     self.check_over_configs(num_train_timesteps=a_,   time_step=a_	)


 def 				UpperCamelCase_     (       self:	Any	):




   '''simple docstring'''



   for t, num_inference_steps in zip([1, 5, 10],   [10, 50, 100]	):
     self.check_over_forward(num_inference_steps=a_,   time_step=a_	)


 def 				UpperCamelCase_     (       self:	Optional[Any]	):




   '''simple docstring'''



   _snake_case							:	Optional[int]			=						self.full_loop()
   _snake_case							:	Optional[int]			=						torch.mean(torch.abs(a_	)	)

   assert abs(result_mean.item() - 2_540_529	) < 10


