



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




				def __init__(  self   :  Dict,   __lowercase   :  Dict,   __lowercase   :  str=13,   __lowercase   :  Tuple=7,   __lowercase   :  str=True,   __lowercase   :  int=True,   __lowercase   :  str=True,   __lowercase   :  Union[str, Any]=True,   __lowercase   :  List[Any]=99,   __lowercase   :  Any=32,   __lowercase   :  List[str]=5,   __lowercase   :  str=4,   __lowercase   :  Optional[Any]=37,   __lowercase   :  Optional[Any]="gelu",   __lowercase   :  Any=0.1,   __lowercase   :  Tuple=0.1,   __lowercase   :  List[str]=512,   __lowercase   :  Optional[Any]=16,   __lowercase   :  Optional[int]=2,   __lowercase   :  Union[str, Any]=0.02,   __lowercase   :  Dict=3,   __lowercase   :  int=4,   __lowercase   :  Optional[Any]=None,   ):
							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   :  Any  ):
							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   :  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=__lowercase,   initializer_range=self.initializer_range,   )




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




				def 							A__ (  self   :  Optional[Any],   __lowercase   :  Tuple,   __lowercase   :  int,   __lowercase   :  Tuple,   __lowercase   :  Optional[int],   __lowercase   :  Any,   __lowercase   :  List[Any],   __lowercase   :  Any  ):
							lowercase__   					=			NystromformerForMaskedLM(config=__lowercase  )
							model.to(__lowercase  )
							model.eval()
							lowercase__   					=			model(__lowercase,   attention_mask=__lowercase,   token_type_ids=__lowercase,   labels=__lowercase  )
							self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.seq_length, self.vocab_size)  )




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




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




				def 							A__ (  self   :  List[str],   __lowercase   :  Any,   __lowercase   :  Dict,   __lowercase   :  int,   __lowercase   :  List[Any],   __lowercase   :  Optional[Any],   __lowercase   :  Tuple,   __lowercase   :  Tuple  ):
							lowercase__   					=			self.num_choices
							lowercase__   					=			NystromformerForMultipleChoice(config=__lowercase  )
							model.to(__lowercase  )
							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(
							    __lowercase,   attention_mask=__lowercase,   token_type_ids=__lowercase,   labels=__lowercase,   )
							self.parent.assertEqual(result.logits.shape,   (self.batch_size, self.num_choices)  )




				def 							A__ (  self   :  Tuple  ):
							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 			_snake_case     (   lowercase__      ,		lowercase__      ,		unittest.TestCase):
				UpperCamelCase__			:   str       						=(
				    (
				        NystromformerModel,
				        NystromformerForMaskedLM,
				        NystromformerForMultipleChoice,
				        NystromformerForQuestionAnswering,
				        NystromformerForSequenceClassification,
				        NystromformerForTokenClassification,
				    )
				    if is_torch_available()
				    else ()
				)
				UpperCamelCase__			:   Tuple       						=(
				    {
				        """feature-extraction""": NystromformerModel,
				        """fill-mask""": NystromformerForMaskedLM,
				        """question-answering""": NystromformerForQuestionAnswering,
				        """text-classification""": NystromformerForSequenceClassification,
				        """token-classification""": NystromformerForTokenClassification,
				        """zero-shot""": NystromformerForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)
				UpperCamelCase__			:   List[Any]       						=False
				UpperCamelCase__			:   Any       						=False




				def 							A__ (  self   :  Optional[Any]  ):
							lowercase__   					=			NystromformerModelTester(self  )
							lowercase__   					=			ConfigTester(self,   config_class=__lowercase,   hidden_size=37  )




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




				def 							A__ (  self   :  int  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__lowercase  )




				def 							A__ (  self   :  Optional[int]  ):
							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(*__lowercase  )




				def 							A__ (  self   :  int  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__lowercase  )




				def 							A__ (  self   :  str  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*__lowercase  )




				def 							A__ (  self   :  int  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*__lowercase  )




				def 							A__ (  self   :  List[Any]  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*__lowercase  )




				def 							A__ (  self   :  Dict  ):
							lowercase__   					=			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__lowercase  )




				@slow
				def 							A__ (  self   :  Any  ):
							for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowercase__   					=			NystromformerModel.from_pretrained(__lowercase  )
										self.assertIsNotNone(__lowercase  )
@require_torch
class 			_snake_case     (   unittest.TestCase):




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

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

							lowercase__   					=			torch.Size((1, 6, 768)  )
							self.assertEqual(output.shape,   __lowercase  )

							lowercase__   					=			torch.tensor(
							    [[[-0.4532, -0.0936, 0.5137], [-0.2676, 0.0628, 0.6186], [-0.3629, -0.1726, 0.4716]]]  )

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




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

							lowercase__   					=			AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"  )
							lowercase__   					=			NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"  )

							lowercase__   					=			tokenizer(__lowercase,   return_tensors="pt"  )

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

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

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









from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_						=       logging.get_logger(__name__)

lowercase_						=       {
    """facebook/dpr-ctx_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-single-nq-base""": (
        """https://huggingface.co/facebook/dpr-reader-single-nq-base/resolve/main/config.json"""
    ),
    """facebook/dpr-ctx_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-ctx_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-question_encoder-multiset-base""": (
        """https://huggingface.co/facebook/dpr-question_encoder-multiset-base/resolve/main/config.json"""
    ),
    """facebook/dpr-reader-multiset-base""": (
        """https://huggingface.co/facebook/dpr-reader-multiset-base/resolve/main/config.json"""
    ),
}
class 			_snake_case     (   lowercase__):
				UpperCamelCase__			:   str       						="""dpr"""




				def __init__(  self   :  int,   __lowercase   :  List[str]=3_0522,   __lowercase   :  Any=768,   __lowercase   :  Union[str, Any]=12,   __lowercase   :  Optional[int]=12,   __lowercase   :  List[Any]=3072,   __lowercase   :  List[str]="gelu",   __lowercase   :  Union[str, Any]=0.1,   __lowercase   :  str=0.1,   __lowercase   :  List[str]=512,   __lowercase   :  Optional[int]=2,   __lowercase   :  Dict=0.02,   __lowercase   :  Any=1e-1_2,   __lowercase   :  List[Any]=0,   __lowercase   :  str="absolute",   __lowercase   :  int = 0,   **__lowercase   :  str,   ):
							super().__init__(pad_token_id=__lowercase,   **__lowercase  )

							lowercase__   					=			vocab_size
							lowercase__   					=			hidden_size
							lowercase__   					=			num_hidden_layers
							lowercase__   					=			num_attention_heads
							lowercase__   					=			hidden_act
							lowercase__   					=			intermediate_size
							lowercase__   					=			hidden_dropout_prob
							lowercase__   					=			attention_probs_dropout_prob
							lowercase__   					=			max_position_embeddings
							lowercase__   					=			type_vocab_size
							lowercase__   					=			initializer_range
							lowercase__   					=			layer_norm_eps
							lowercase__   					=			projection_dim
							lowercase__   					=			position_embedding_type




