


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




		def __init__(		self   ,       UpperCamelCase_   ,       UpperCamelCase_=13   ,       UpperCamelCase_=7   ,       UpperCamelCase_=True   ,       UpperCamelCase_=True   ,       UpperCamelCase_=True   ,       UpperCamelCase_=True   ,       UpperCamelCase_=99   ,       UpperCamelCase_=32   ,       UpperCamelCase_=5   ,       UpperCamelCase_=4   ,       UpperCamelCase_=37   ,       UpperCamelCase_="gelu"   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=512   ,       UpperCamelCase_=16   ,       UpperCamelCase_=2   ,       UpperCamelCase_=0.02   ,       UpperCamelCase_=3   ,       UpperCamelCase_=4   ,       UpperCamelCase_=None   ,       ):

							'''simple docstring'''


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




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[str]							 =			ids_tensor([self.batch_size, self.seq_length]   ,       self.vocab_size							)

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

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

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

							UpperCamelCase__       :Dict							 =			self.get_config()

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




		def    lowerCAmelCase__    (		self							):

							'''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=_lowerCamelCase   ,       initializer_range=self.initializer_range   ,       )




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			NystromformerModel(config=_lowerCamelCase							)
							model.to(_lowerCamelCase							)
							model.eval()
							UpperCamelCase__       :List[Any]							 =			model(_lowerCamelCase   ,       attention_mask=_lowerCamelCase   ,       token_type_ids=_lowerCamelCase							)
							UpperCamelCase__       :Tuple							 =			model(_lowerCamelCase   ,       token_type_ids=_lowerCamelCase							)
							UpperCamelCase__       :Dict							 =			model(_lowerCamelCase							)
							self.parent.assertEqual(result.last_hidden_state.shape   ,       (self.batch_size, self.seq_length, self.hidden_size)							)




		def    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[int]							 =			NystromformerForMaskedLM(config=_lowerCamelCase							)
							model.to(_lowerCamelCase							)
							model.eval()
							UpperCamelCase__       :Tuple							 =			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    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[Any]							 =			NystromformerForQuestionAnswering(config=_lowerCamelCase							)
							model.to(_lowerCamelCase							)
							model.eval()
							UpperCamelCase__       :List[str]							 =			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    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.num_labels
							UpperCamelCase__       :Tuple							 =			NystromformerForSequenceClassification(_lowerCamelCase							)
							model.to(_lowerCamelCase							)
							model.eval()
							UpperCamelCase__       :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    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			self.num_labels
							UpperCamelCase__       :int							 =			NystromformerForTokenClassification(config=_lowerCamelCase							)
							model.to(_lowerCamelCase							)
							model.eval()
							UpperCamelCase__       :Optional[int]							 =			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    lowerCAmelCase__    (		self   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_   ,       UpperCamelCase_							):

							'''simple docstring'''


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




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			self.prepare_config_and_inputs()
							(
							    UpperCamelCase__
							)       :Tuple							 =			config_and_inputs
							UpperCamelCase__       :Optional[int]							 =			{'''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 ):
		_a   			=    (
		    (
		        NystromformerModel,
		        NystromformerForMaskedLM,
		        NystromformerForMultipleChoice,
		        NystromformerForQuestionAnswering,
		        NystromformerForSequenceClassification,
		        NystromformerForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		_a   			=    (
		    {
		        'feature-extraction': NystromformerModel,
		        'fill-mask': NystromformerForMaskedLM,
		        'question-answering': NystromformerForQuestionAnswering,
		        'text-classification': NystromformerForSequenceClassification,
		        'token-classification': NystromformerForTokenClassification,
		        'zero-shot': NystromformerForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		_a   			=    False
		_a   			=    False




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			NystromformerModelTester(self							)
							UpperCamelCase__       :int							 =			ConfigTester(self   ,       config_class=_lowerCamelCase   ,       hidden_size=37							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							self.config_tester.run_common_tests()




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*_lowerCamelCase							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


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




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Union[str, Any]							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Optional[int]							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_multiple_choice(*_lowerCamelCase							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :Tuple							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_question_answering(*_lowerCamelCase							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[str]							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_sequence_classification(*_lowerCamelCase							)




		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[str]							 =			self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase							)




		@slow
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
												UpperCamelCase__       :Optional[Any]							 =			NystromformerModel.from_pretrained(_lowerCamelCase							)
												self.assertIsNotNone(_lowerCamelCase							)



@require_torch
class        lowercase     (							unittest.TestCase ):




		@slow
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			NystromformerModel.from_pretrained('''uw-madison/nystromformer-512'''							)
							UpperCamelCase__       :Optional[int]							 =			torch.tensor([[0, 1, 2, 3, 4, 5]]							)

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

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

							UpperCamelCase__       :Union[str, Any]							 =			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]   ,       _lowerCamelCase   ,       atol=1e-4							)							)




		@slow
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			'''the [MASK] of Belgium is Brussels'''

							UpperCamelCase__       :Union[str, Any]							 =			AutoTokenizer.from_pretrained('''uw-madison/nystromformer-512'''							)
							UpperCamelCase__       :Dict							 =			NystromformerForMaskedLM.from_pretrained('''uw-madison/nystromformer-512'''							)

							UpperCamelCase__       :Optional[int]							 =			tokenizer(_lowerCamelCase   ,       return_tensors='''pt'''							)

							with torch.no_grad():
												UpperCamelCase__       :Optional[int]							 =			model(encoding.input_ids							).logits

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

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



'''simple docstring'''
from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__snake_case											=					logging.get_logger(__name__)

__snake_case											=					{
    '''distilbert-base-uncased''': '''https://huggingface.co/distilbert-base-uncased/resolve/main/config.json''',
    '''distilbert-base-uncased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-cased''': '''https://huggingface.co/distilbert-base-cased/resolve/main/config.json''',
    '''distilbert-base-cased-distilled-squad''': (
        '''https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json'''
    ),
    '''distilbert-base-german-cased''': '''https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json''',
    '''distilbert-base-multilingual-cased''': (
        '''https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json'''
    ),
    '''distilbert-base-uncased-finetuned-sst-2-english''': (
        '''https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json'''
    ),
}
class        lowercase     (							A__ ):

		"""simple docstring"""

		_a   			=    'distilbert'
		_a   			=    {
		    'hidden_size': 'dim',
		    'num_attention_heads': 'n_heads',
		    'num_hidden_layers': 'n_layers',
		}




		def __init__(		self   ,       UpperCamelCase_=30522   ,       UpperCamelCase_=512   ,       UpperCamelCase_=False   ,       UpperCamelCase_=6   ,       UpperCamelCase_=12   ,       UpperCamelCase_=768   ,       UpperCamelCase_=4 * 768   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_="gelu"   ,       UpperCamelCase_=0.02   ,       UpperCamelCase_=0.1   ,       UpperCamelCase_=0.2   ,       UpperCamelCase_=0   ,       **UpperCamelCase_   ,       ):

							'''simple docstring'''


							UpperCamelCase__       :List[Any]							 =			vocab_size
							UpperCamelCase__       :Dict							 =			max_position_embeddings
							UpperCamelCase__       :str							 =			sinusoidal_pos_embds
							UpperCamelCase__       :Any							 =			n_layers
							UpperCamelCase__       :str							 =			n_heads
							UpperCamelCase__       :Tuple							 =			dim
							UpperCamelCase__       :str							 =			hidden_dim
							UpperCamelCase__       :Dict							 =			dropout
							UpperCamelCase__       :int							 =			attention_dropout
							UpperCamelCase__       :Optional[Any]							 =			activation
							UpperCamelCase__       :Optional[int]							 =			initializer_range
							UpperCamelCase__       :Union[str, Any]							 =			qa_dropout
							UpperCamelCase__       :Dict							 =			seq_classif_dropout
							super().__init__(**UpperCamelCase_   ,       pad_token_id=UpperCamelCase_							)



class        lowercase     (							A__ ):

		"""simple docstring"""

		@property
		def    lowerCAmelCase__    (		self							):

							'''simple docstring'''


							if self.task == "multiple-choice":
												UpperCamelCase__       :str							 =			{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
							else:
												UpperCamelCase__       :str							 =			{0: '''batch''', 1: '''sequence'''}
							return OrderedDict(
							    [
							        ('''input_ids''', dynamic_axis),
							        ('''attention_mask''', dynamic_axis),
							    ]							)