





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






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






							def        __a  (  self )      ->					str:
								a :						str      =							ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size )

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

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

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

								a :						str      =							self.get_config()

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






							def        __a  (  self )      ->					Dict:
								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=lowerCAmelCase__ ,   initializer_range=self.initializer_range ,   )






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					Tuple:
								a :						Any      =							NystromformerModel(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						List[Any]      =							model(lowerCAmelCase__ ,   attention_mask=lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ )
								a :						int      =							model(lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ )
								a :						Any      =							model(lowerCAmelCase__ )
								self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size) )






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					Dict:
								a :						List[str]      =							NystromformerForMaskedLM(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						Dict      =							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 ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					Optional[Any]:
								a :						List[Any]      =							NystromformerForQuestionAnswering(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						List[Any]      =							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 ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					List[str]:
								a :						List[Any]      =							self.num_labels
								a :						Tuple      =							NystromformerForSequenceClassification(lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						Optional[int]      =							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 ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					List[str]:
								a :						Dict      =							self.num_labels
								a :						str      =							NystromformerForTokenClassification(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						str      =							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 ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					Union[str, Any]:
								a :						int      =							self.num_choices
								a :						List[Any]      =							NystromformerForMultipleChoice(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						Union[str, Any]      =							input_ids.unsqueeze(1 ).expand(-1 ,   self.num_choices ,   -1 ).contiguous()
								a :						List[Any]      =							token_type_ids.unsqueeze(1 ).expand(-1 ,   self.num_choices ,   -1 ).contiguous()
								a :						Optional[Any]      =							input_mask.unsqueeze(1 ).expand(-1 ,   self.num_choices ,   -1 ).contiguous()
								a :						List[str]      =							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 )      ->					Any:
								a :						List[Any]      =							self.prepare_config_and_inputs()
								(
								    (
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			
								) :						List[Any]      =							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 			__UpperCamelCase							(		a__						,			a__						,			unittest.TestCase		):
							lowerCamelCase  :   List[Any]    =(
							    (
							        NystromformerModel,
							        NystromformerForMaskedLM,
							        NystromformerForMultipleChoice,
							        NystromformerForQuestionAnswering,
							        NystromformerForSequenceClassification,
							        NystromformerForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							lowerCamelCase  :   str    =(
							    {
							        """feature-extraction""": NystromformerModel,
							        """fill-mask""": NystromformerForMaskedLM,
							        """question-answering""": NystromformerForQuestionAnswering,
							        """text-classification""": NystromformerForSequenceClassification,
							        """token-classification""": NystromformerForTokenClassification,
							        """zero-shot""": NystromformerForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowerCamelCase  :   Tuple    =False
							lowerCamelCase  :   Optional[int]    =False






							def        __a  (  self )      ->					Optional[Any]:
								a :						Union[str, Any]      =							NystromformerModelTester(self )
								a :						List[str]      =							ConfigTester(self ,   config_class=lowerCAmelCase__ ,   hidden_size=37 )






							def        __a  (  self )      ->					Dict:
								self.config_tester.run_common_tests()






							def        __a  (  self )      ->					List[str]:
								a :						Optional[Any]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowerCAmelCase__ )






							def        __a  (  self )      ->					Tuple:
								a :						Dict      =							self.model_tester.prepare_config_and_inputs()
								for type in ["absolute", "relative_key", "relative_key_query"]:
									a :						Optional[int]      =							type
									self.model_tester.create_and_check_model(*lowerCAmelCase__ )






							def        __a  (  self )      ->					str:
								a :						Dict      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__ )






							def        __a  (  self )      ->					Any:
								a :						Union[str, Any]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__ )






							def        __a  (  self )      ->					str:
								a :						Any      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__ )






							def        __a  (  self )      ->					Optional[int]:
								a :						Optional[int]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__ )






							def        __a  (  self )      ->					Tuple:
								a :						Union[str, Any]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__ )






							@slow
							def        __a  (  self )      ->					Dict:
								for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									a :						Optional[Any]      =							NystromformerModel.from_pretrained(lowerCAmelCase__ )
									self.assertIsNotNone(lowerCAmelCase__ )


@require_torch
class 			__UpperCamelCase							(		unittest.TestCase		):






							@slow
							def        __a  (  self )      ->					Union[str, Any]:
								a :						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 :						Dict      =							model(lowerCAmelCase__ )[0]

								a :						Dict      =							torch.Size((1, 6, 768) )
								self.assertEqual(output.shape ,   lowerCAmelCase__ )

								a :						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] ,   lowerCAmelCase__ ,   atol=1E-4 ) )






							@slow
							def        __a  (  self )      ->					Tuple:
								a :						List[Any]      =							"the [MASK] of Belgium is Brussels"

								a :						Dict      =							AutoTokenizer.from_pretrained("uw-madison/nystromformer-512" )
								a :						Dict      =							NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512" )

								a :						Optional[Any]      =							tokenizer(lowerCAmelCase__ ,   return_tensors="pt" )

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

								a :						Union[str, Any]      =							token_logits[:, 2, :].argmax(-1 )[0]

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









"""simple docstring"""




from ..utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
)


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_pt_objects import *  # noqa F403
else:
		from .scheduling_consistency_models import CMStochasticIterativeScheduler
		from .scheduling_ddim import DDIMScheduler
		from .scheduling_ddim_inverse import DDIMInverseScheduler
		from .scheduling_ddim_parallel import DDIMParallelScheduler
		from .scheduling_ddpm import DDPMScheduler
		from .scheduling_ddpm_parallel import DDPMParallelScheduler
		from .scheduling_deis_multistep import DEISMultistepScheduler
		from .scheduling_dpmsolver_multistep import DPMSolverMultistepScheduler
		from .scheduling_dpmsolver_multistep_inverse import DPMSolverMultistepInverseScheduler
		from .scheduling_dpmsolver_singlestep import DPMSolverSinglestepScheduler
		from .scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
		from .scheduling_euler_discrete import EulerDiscreteScheduler
		from .scheduling_heun_discrete import HeunDiscreteScheduler
		from .scheduling_ipndm import IPNDMScheduler
		from .scheduling_k_dpm_2_ancestral_discrete import KDPMaAncestralDiscreteScheduler
		from .scheduling_k_dpm_2_discrete import KDPMaDiscreteScheduler
		from .scheduling_karras_ve import KarrasVeScheduler
		from .scheduling_pndm import PNDMScheduler
		from .scheduling_repaint import RePaintScheduler
		from .scheduling_sde_ve import ScoreSdeVeScheduler
		from .scheduling_sde_vp import ScoreSdeVpScheduler
		from .scheduling_unclip import UnCLIPScheduler
		from .scheduling_unipc_multistep import UniPCMultistepScheduler
		from .scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin
		from .scheduling_vq_diffusion import VQDiffusionScheduler

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_flax_objects import *  # noqa F403
else:
		from .scheduling_ddim_flax import FlaxDDIMScheduler
		from .scheduling_ddpm_flax import FlaxDDPMScheduler
		from .scheduling_dpmsolver_multistep_flax import FlaxDPMSolverMultistepScheduler
		from .scheduling_karras_ve_flax import FlaxKarrasVeScheduler
		from .scheduling_lms_discrete_flax import FlaxLMSDiscreteScheduler
		from .scheduling_pndm_flax import FlaxPNDMScheduler
		from .scheduling_sde_ve_flax import FlaxScoreSdeVeScheduler
		from .scheduling_utils_flax import (
		    FlaxKarrasDiffusionSchedulers,
		    FlaxSchedulerMixin,
		    FlaxSchedulerOutput,
		    broadcast_to_shape_from_left,
		)


try:
		if not (is_torch_available() and is_scipy_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
		from .scheduling_lms_discrete import LMSDiscreteScheduler

try:
		if not (is_torch_available() and is_torchsde_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ..utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
		from .scheduling_dpmsolver_sde import DPMSolverSDEScheduler


