




import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


snake_case_ : int    		=       "bert-base-cased"
snake_case_ : List[Any]    		=       "google/pegasus-xsum"
snake_case_ : Optional[Any]    		=       [" Sam ate lunch today.", "Sams lunch ingredients."]
snake_case_ : Any    		=       ["A very interesting story about what I ate for lunch.", "Avocado, celery, turkey, coffee"]
snake_case_ : List[str]    		=       "patrickvonplaten/t5-tiny-random"
snake_case_ : Tuple    		=       "sshleifer/bart-tiny-random"
snake_case_ : List[str]    		=       "sshleifer/tiny-mbart"
snake_case_ : List[Any]    		=       "sshleifer/tiny-marian-en-de"


def    A						(__A   :    Path   ,				__A   :    list       )   ->      Tuple:







     """simple docstring"""


     UpperCAmelCase_          =							'''\n'''.join(__A       )
     Path(__A       ).open('''w'''       ).writelines(__A       )


def    A						(__A   :    str       )   ->      Union[str, Any]:







     """simple docstring"""


     for split in ["train", "val", "test"]:
          _dump_articles(os.path.join(__A   ,				F"""{split}.source"""       )   ,				__A       )
          _dump_articles(os.path.join(__A   ,				F"""{split}.target"""       )   ,				__A       )
     return tmp_dir







class __snake_case			(			a		):







   @parameterized.expand(
       [
           MBART_TINY,
           MARIAN_TINY,
           T5_TINY,
           BART_TINY,
           PEGASUS_XSUM,
       ]					,						)
   @slow
   def lowerCamelCase						(      self		:						List[str]					,						_snake_case		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							AutoTokenizer.from_pretrained(_snake_case)
        UpperCAmelCase_          =							make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        UpperCAmelCase_          =							max(len(tokenizer.encode(_snake_case)) for a in ARTICLES)
        UpperCAmelCase_          =							max(len(tokenizer.encode(_snake_case)) for a in SUMMARIES)
        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							8
        assert max_len_target > max_src_len  # Will be truncated
        assert max_len_source > max_src_len  # Will be truncated
        UpperCAmelCase_     ,       UpperCAmelCase_          =							'''ro_RO''', '''de_DE'''  # ignored for all but mbart, but never causes error.
        UpperCAmelCase_          =							SeqaSeqDataset(
            _snake_case					,						data_dir=_snake_case					,						type_path='''train'''					,						max_source_length=_snake_case					,						max_target_length=_snake_case					,						src_lang=_snake_case					,						tgt_lang=_snake_case					,						)
        UpperCAmelCase_          =							DataLoader(_snake_case					,						batch_size=2					,						collate_fn=train_dataset.collate_fn)
        for batch in dataloader:
             assert isinstance(_snake_case					,						_snake_case)
             assert batch["attention_mask"].shape == batch["input_ids"].shape
             # show that articles were trimmed.
             assert batch["input_ids"].shape[1] == max_src_len
             # show that targets are the same len
             assert batch["labels"].shape[1] == max_tgt_len
             if tok_name != MBART_TINY:
                  continue
             # check language codes in correct place
             UpperCAmelCase_          =							shift_tokens_right(batch['''labels''']					,						tokenizer.pad_token_id)
             assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
             assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
             assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
             assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

             break  # No need to test every batch







   @parameterized.expand([BART_TINY, BERT_BASE_CASED])
   def lowerCamelCase						(      self		:						Optional[int]					,						_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							AutoTokenizer.from_pretrained(_snake_case)
        UpperCAmelCase_          =							make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())
        UpperCAmelCase_          =							max(len(tokenizer.encode(_snake_case)) for a in ARTICLES)
        UpperCAmelCase_          =							max(len(tokenizer.encode(_snake_case)) for a in SUMMARIES)
        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							LegacySeqaSeqDataset(
            _snake_case					,						data_dir=_snake_case					,						type_path='''train'''					,						max_source_length=20					,						max_target_length=_snake_case					,						)
        UpperCAmelCase_          =							DataLoader(_snake_case					,						batch_size=2					,						collate_fn=train_dataset.collate_fn)
        for batch in dataloader:
             assert batch["attention_mask"].shape == batch["input_ids"].shape
             # show that articles were trimmed.
             assert batch["input_ids"].shape[1] == max_len_source
             assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
             # show that targets were truncated
             assert batch["labels"].shape[1] == trunc_target  # Truncated
             assert max_len_target > trunc_target  # Truncated
             break  # No need to test every batch







   def lowerCamelCase						(      self		:						Optional[int]):

        """simple docstring"""


        UpperCAmelCase_          =							AutoTokenizer.from_pretrained('''facebook/mbart-large-cc25''')

        UpperCAmelCase_          =							Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
        UpperCAmelCase_          =							tmp_dir.joinpath('''train.source''').open().readlines()
        UpperCAmelCase_          =							Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()))
        pack_data_dir(_snake_case					,						_snake_case					,						128					,						_snake_case)
        UpperCAmelCase_          =							{x.name for x in tmp_dir.iterdir()}
        UpperCAmelCase_          =							{x.name for x in save_dir.iterdir()}
        UpperCAmelCase_          =							save_dir.joinpath('''train.source''').open().readlines()
        # orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
        # desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
        assert len(_snake_case) < len(_snake_case)
        assert len(_snake_case) == 1
        assert len(packed_examples[0]) == sum(len(_snake_case) for x in orig_examples)
        assert orig_paths == new_paths







   @pytest.mark.skipif(not FAIRSEQ_AVAILABLE					,						reason='''This test requires fairseq''')
   def lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        if not FAIRSEQ_AVAILABLE:
             return
        UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							self._get_dataset(max_len=64)
        UpperCAmelCase_          =							64
        UpperCAmelCase_          =							ds.make_dynamic_sampler(_snake_case					,						required_batch_size_multiple=_snake_case)
        UpperCAmelCase_          =							[len(_snake_case) for x in batch_sampler]
        assert len(set(_snake_case)) > 1  # it's not dynamic batch size if every batch is the same length
        assert sum(_snake_case) == len(_snake_case)  # no dropped or added examples
        UpperCAmelCase_          =							DataLoader(_snake_case					,						batch_sampler=_snake_case					,						collate_fn=ds.collate_fn					,						num_workers=2)
        UpperCAmelCase_          =							[]
        UpperCAmelCase_          =							[]
        for batch in data_loader:
             UpperCAmelCase_          =							batch['''input_ids'''].shape
             UpperCAmelCase_          =							src_shape[0]
             assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
             UpperCAmelCase_          =							np.product(batch['''input_ids'''].shape)
             num_src_per_batch.append(_snake_case)
             if num_src_tokens > (max_tokens * 1.1):
                  failures.append(_snake_case)
        assert num_src_per_batch[0] == max(_snake_case)
        if failures:
             raise AssertionError(F"""too many tokens in {len(_snake_case)} batches""")







   def lowerCamelCase						(      self		:						int):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							self._get_dataset(max_len=512)
        UpperCAmelCase_          =							2
        UpperCAmelCase_          =							ds.make_sortish_sampler(_snake_case					,						shuffle=_snake_case)

        UpperCAmelCase_          =							DataLoader(_snake_case					,						batch_size=_snake_case					,						collate_fn=ds.collate_fn					,						num_workers=2)
        UpperCAmelCase_          =							DataLoader(_snake_case					,						batch_size=_snake_case					,						collate_fn=ds.collate_fn					,						num_workers=2					,						sampler=_snake_case)

        UpperCAmelCase_          =							tokenizer.pad_token_id

        def count_pad_tokens(_snake_case		:						List[Any]					,						_snake_case		:						Tuple="input_ids"):
             return [batch[k].eq(_snake_case).sum().item() for batch in data_loader]

        assert sum(count_pad_tokens(_snake_case					,						k='''labels''')) < sum(count_pad_tokens(_snake_case					,						k='''labels'''))
        assert sum(count_pad_tokens(_snake_case)) < sum(count_pad_tokens(_snake_case))
        assert len(_snake_case) == len(_snake_case)







   def lowerCamelCase						(      self		:						Optional[Any]					,						_snake_case		:						Dict=1000					,						_snake_case		:						str=128):

        """simple docstring"""


        if os.getenv('''USE_REAL_DATA'''					,						_snake_case):
             UpperCAmelCase_          =							'''examples/seq2seq/wmt_en_ro'''
             UpperCAmelCase_          =							max_len * 2 * 64
             if not Path(_snake_case).joinpath('''train.len''').exists():
                  save_len_file(_snake_case					,						_snake_case)
        else:
             UpperCAmelCase_          =							'''examples/seq2seq/test_data/wmt_en_ro'''
             UpperCAmelCase_          =							max_len * 4
             save_len_file(_snake_case					,						_snake_case)

        UpperCAmelCase_          =							AutoTokenizer.from_pretrained(_snake_case)
        UpperCAmelCase_          =							SeqaSeqDataset(
            _snake_case					,						data_dir=_snake_case					,						type_path='''train'''					,						max_source_length=_snake_case					,						max_target_length=_snake_case					,						n_obs=_snake_case					,						)
        return ds, max_tokens, tokenizer







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_     ,       UpperCAmelCase_     ,       UpperCAmelCase_          =							self._get_dataset()
        UpperCAmelCase_          =							set(DistributedSortishSampler(_snake_case					,						256					,						num_replicas=2					,						rank=0					,						add_extra_examples=_snake_case))
        UpperCAmelCase_          =							set(DistributedSortishSampler(_snake_case					,						256					,						num_replicas=2					,						rank=1					,						add_extra_examples=_snake_case))
        assert idsa.intersection(_snake_case) == set()







   @parameterized.expand(
       [
           MBART_TINY,
           MARIAN_TINY,
           T5_TINY,
           BART_TINY,
           PEGASUS_XSUM,
       ]					,						)
   def lowerCamelCase						(      self		:						List[str]					,						_snake_case		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							AutoTokenizer.from_pretrained(_snake_case					,						use_fast=_snake_case)
        if tok_name == MBART_TINY:
             UpperCAmelCase_          =							SeqaSeqDataset(
                 _snake_case					,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())					,						type_path='''train'''					,						max_source_length=4					,						max_target_length=8					,						src_lang='''EN'''					,						tgt_lang='''FR'''					,						)
             UpperCAmelCase_          =							train_dataset.dataset_kwargs
             assert "src_lang" in kwargs and "tgt_lang" in kwargs
        else:
             UpperCAmelCase_          =							SeqaSeqDataset(
                 _snake_case					,						data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir())					,						type_path='''train'''					,						max_source_length=4					,						max_target_length=8					,						)
             UpperCAmelCase_          =							train_dataset.dataset_kwargs
             assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
             assert len(_snake_case) == 1 if tok_name == BART_TINY else len(_snake_case) == 0







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







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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										lowercase_					:			int										=       ids_tensor([self.batch_size, self.seq_length]				,			self.vocab_size    )

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

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

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

										lowercase_					:			Tuple										=       self.get_config()

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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										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 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Any				,			lowercase_				:  List[str]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  Tuple				,			lowercase_				:  Tuple				,			lowercase_				:  Optional[int]    ):
										lowercase_					:			Optional[Any]										=       NystromformerModel(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Optional[int]										=       model(lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_    )
										lowercase_					:			Optional[Any]										=       model(lowercase_				,			token_type_ids=lowercase_    )
										lowercase_					:			Union[str, Any]										=       model(lowercase_    )
										self.parent.assertEqual(result.last_hidden_state.shape				,			(self.batch_size, self.seq_length, self.hidden_size)    )







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







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]				,			lowercase_				:  List[str]				,			lowercase_				:  int				,			lowercase_				:  List[str]				,			lowercase_				:  Dict				,			lowercase_				:  Tuple				,			lowercase_				:  Union[str, Any]				,			lowercase_				:  List[str]    ):
										lowercase_					:			int										=       self.num_labels
										lowercase_					:			int										=       NystromformerForTokenClassification(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			Tuple										=       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 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]				,			lowercase_				:  Tuple				,			lowercase_				:  Dict				,			lowercase_				:  int				,			lowercase_				:  List[str]				,			lowercase_				:  Optional[Any]				,			lowercase_				:  str				,			lowercase_				:  Union[str, Any]    ):
										lowercase_					:			str										=       self.num_choices
										lowercase_					:			Union[str, Any]										=       NystromformerForMultipleChoice(config=lowercase_    )
										model.to(lowercase_    )
										model.eval()
										lowercase_					:			str										=       input_ids.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Tuple										=       token_type_ids.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Dict										=       input_mask.unsqueeze(1    ).expand(-1				,			self.num_choices				,			-1    ).contiguous()
										lowercase_					:			Union[str, Any]										=       model(
										    lowercase_				,			attention_mask=lowercase_				,			token_type_ids=lowercase_				,			labels=lowercase_				,			)
										self.parent.assertEqual(result.logits.shape				,			(self.batch_size, self.num_choices)    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			Union[str, Any]										=       self.prepare_config_and_inputs()
										(
										    (
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						(
										    lowercase_
										)						,						
										)					:			Tuple										=       config_and_inputs
										lowercase_					:			Optional[Any]										=       {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
										return config, inputs_dict





@require_torch
class     __magic_name__						(    _UpperCAmelCase,     _UpperCAmelCase,     unittest.TestCase):
			UpperCamelCase__				=       (
			    (
			        NystromformerModel,
			        NystromformerForMaskedLM,
			        NystromformerForMultipleChoice,
			        NystromformerForQuestionAnswering,
			        NystromformerForSequenceClassification,
			        NystromformerForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			UpperCamelCase__				=       (
			    {
			        '''feature-extraction''': NystromformerModel,
			        '''fill-mask''': NystromformerForMaskedLM,
			        '''question-answering''': NystromformerForQuestionAnswering,
			        '''text-classification''': NystromformerForSequenceClassification,
			        '''token-classification''': NystromformerForTokenClassification,
			        '''zero-shot''': NystromformerForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			UpperCamelCase__				=       False
			UpperCamelCase__				=       False







			def 	SCREAMING_SNAKE_CASE_     (	self				:  int    ):
										lowercase_					:			Any										=       NystromformerModelTester(self    )
										lowercase_					:			Optional[Any]										=       ConfigTester(self				,			config_class=lowercase_				,			hidden_size=37    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Tuple    ):
										self.config_tester.run_common_tests()







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Dict    ):
										lowercase_					:			Dict										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Union[str, Any]										=       self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																	lowercase_					:			int										=       type
																	self.model_tester.create_and_check_model(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			Tuple										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			int										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_multiple_choice(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[Any]    ):
										lowercase_					:			Union[str, Any]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_question_answering(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Any										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_sequence_classification(*lowercase_    )







			def 	SCREAMING_SNAKE_CASE_     (	self				:  List[Any]    ):
										lowercase_					:			Optional[int]										=       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*lowercase_    )







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	lowercase_					:			List[Any]										=       NystromformerModel.from_pretrained(lowercase_    )
																	self.assertIsNotNone(lowercase_    )





@require_torch
class     __magic_name__						(    unittest.TestCase):







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Union[str, Any]    ):
										lowercase_					:			List[str]										=       NystromformerModel.from_pretrained("""uw-madison/nystromformer-512"""    )
										lowercase_					:			List[str]										=       torch.tensor([[0, 1, 2, 3, 4, 5]]    )

										with torch.no_grad():
																	lowercase_					:			Tuple										=       model(lowercase_    )[0]

										lowercase_					:			Tuple										=       torch.Size((1, 6, 768)    )
										self.assertEqual(output.shape				,			lowercase_    )

										lowercase_					:			Dict										=       torch.tensor(
										    [[[-0.45_32, -0.09_36, 0.51_37], [-0.26_76, 0.06_28, 0.61_86], [-0.36_29, -0.17_26, 0.47_16]]]    )

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







			@slow
			def 	SCREAMING_SNAKE_CASE_     (	self				:  Optional[int]    ):
										lowercase_					:			Optional[int]										=       """the [MASK] of Belgium is Brussels"""

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

										lowercase_					:			str										=       tokenizer(lowercase_				,			return_tensors="""pt"""    )

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

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

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





