





import shutil
import tempfile
import unittest

from transformers import SPIECE_UNDERLINE, BatchEncoding, MBartaaTokenizer, MBartaaTokenizerFast, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)

from ...test_tokenization_common import TokenizerTesterMixin


UpperCAmelCase     :     str							=	get_tests_dir("""fixtures/test_sentencepiece.model""")

if is_torch_available():
     from transformers.models.mbart.modeling_mbart import shift_tokens_right

UpperCAmelCase     :     Optional[int]							=	250004
UpperCAmelCase     :     Any							=	250020







@require_sentencepiece
@require_tokenizers
class     __lowerCAmelCase	(   UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : Optional[int]            =      MBartaaTokenizer
      _lowercase    : Dict            =      MBartaaTokenizerFast
      _lowercase    : Dict            =      True
      _lowercase    : str            =      True







      def        _lowercase  (   self							)       ->						Any:



        '''simple docstring'''




        super().setUp()

        # We have a SentencePiece fixture for testing
        a__     : str  =MBartaaTokenizer(lowerCAmelCase__		,				src_lang="en_XX"		,				tgt_lang="ro_RO"		,				keep_accents=lowerCAmelCase__							)
        tokenizer.save_pretrained(self.tmpdirname							)







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : Optional[Any]  ="<s>"
        a__     : Tuple  =0

        self.assertEqual(self.get_tokenizer()._convert_token_to_id(lowerCAmelCase__							)		,				lowerCAmelCase__							)
        self.assertEqual(self.get_tokenizer()._convert_id_to_token(lowerCAmelCase__							)		,				lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        a__     : List[Any]  =list(self.get_tokenizer().get_vocab().keys()							)

        self.assertEqual(vocab_keys[0]		,				"<s>"							)
        self.assertEqual(vocab_keys[1]		,				"<pad>"							)
        self.assertEqual(vocab_keys[-1]		,				"<mask>"							)
        self.assertEqual(len(lowerCAmelCase__							)		,				1_0_5_4							)







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




        self.assertEqual(self.get_tokenizer().vocab_size		,				1_0_5_4							)







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        a__     : Union[str, Any]  =MBartaaTokenizer(lowerCAmelCase__		,				src_lang="en_XX"		,				tgt_lang="ro_RO"		,				keep_accents=lowerCAmelCase__							)

        a__     : int  =tokenizer.tokenize("This is a test"							)
        self.assertListEqual(lowerCAmelCase__		,				["▁This", "▁is", "▁a", "▁t", "est"]							)

        self.assertListEqual(
            tokenizer.convert_tokens_to_ids(lowerCAmelCase__							)		,				[value + tokenizer.fairseq_offset for value in [2_8_5, 4_6, 1_0, 1_7_0, 3_8_2]]		,				)

        a__     : List[str]  =tokenizer.tokenize("I was born in 92000, and this is falsé."							)
        self.assertListEqual(
            lowerCAmelCase__		,				[SPIECE_UNDERLINE + "I", SPIECE_UNDERLINE + "was", SPIECE_UNDERLINE + "b", "or", "n", SPIECE_UNDERLINE + "in", SPIECE_UNDERLINE + "", "9", "2", "0", "0", "0", ",", SPIECE_UNDERLINE + "and", SPIECE_UNDERLINE + "this", SPIECE_UNDERLINE + "is", SPIECE_UNDERLINE + "f", "al", "s", "é", "."]		,				)
        a__     : Optional[Any]  =tokenizer.convert_tokens_to_ids(lowerCAmelCase__							)
        self.assertListEqual(
            lowerCAmelCase__		,				[
                value + tokenizer.fairseq_offset
                for value in [8, 2_1, 8_4, 5_5, 2_4, 1_9, 7, 2, 6_0_2, 3_4_7, 3_4_7, 3_4_7, 3, 1_2, 6_6, 4_6, 7_2, 8_0, 6, 2, 4]
            ]		,				)

        a__     : Dict  =tokenizer.convert_ids_to_tokens(lowerCAmelCase__							)
        self.assertListEqual(
            lowerCAmelCase__		,				[SPIECE_UNDERLINE + "I", SPIECE_UNDERLINE + "was", SPIECE_UNDERLINE + "b", "or", "n", SPIECE_UNDERLINE + "in", SPIECE_UNDERLINE + "", "<unk>", "2", "0", "0", "0", ",", SPIECE_UNDERLINE + "and", SPIECE_UNDERLINE + "this", SPIECE_UNDERLINE + "is", SPIECE_UNDERLINE + "f", "al", "s", "<unk>", "."]		,				)







      @slow
      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : List[Any]  ={"input_ids": [[2_5_0_0_0_4, 1_1_0_6_2, 8_2_7_7_2, 7, 1_5, 8_2_7_7_2, 5_3_8, 5_1_5_2_9, 2_3_7, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 2_1_5_1_7_5, 1_3_1_4, 1_3_6, 1_7_1_9_8, 1_2_9_0, 2_0_6, 9, 5_6_3_5_9, 4_2, 1_2_2_0_0_9, 9, 1_6_4_6_6, 1_6, 8_7_3_4_4, 4_5_3_7, 9, 4_7_1_7, 7_8_3_8_1, 6, 1_5_9_9_5_8, 7, 1_5, 2_4_4_8_0, 6_1_8, 4, 5_2_7, 2_2_6_9_3, 5_4_2_8, 4, 2_7_7_7, 2_4_4_8_0, 9_8_7_4, 4, 4_3_5_2_3, 5_9_4, 4, 8_0_3, 1_8_3_9_2, 3_3_1_8_9, 1_8, 4, 4_3_5_2_3, 2_4_4_4_7, 1_2_3_9_9, 1_0_0, 2_4_9_5_5, 8_3_6_5_8, 9_6_2_6, 1_4_4_0_5_7, 1_5, 8_3_9, 2_2_3_3_5, 1_6, 1_3_6, 2_4_9_5_5, 8_3_6_5_8, 8_3_4_7_9, 1_5, 3_9_1_0_2, 7_2_4, 1_6, 6_7_8, 6_4_5, 2_7_8_9, 1_3_2_8, 4_5_8_9, 4_2, 1_2_2_0_0_9, 1_1_5_7_7_4, 2_3, 8_0_5, 1_3_2_8, 4_6_8_7_6, 7, 1_3_6, 5_3_8_9_4, 1_9_4_0, 4_2_2_2_7, 4_1_1_5_9, 1_7_7_2_1, 8_2_3, 4_2_5, 4, 2_7_5_1_2, 9_8_7_2_2, 2_0_6, 1_3_6, 5_5_3_1, 4_9_7_0, 9_1_9, 1_7_3_3_6, 5, 2], [2_5_0_0_0_4, 2_0_0_8_0, 6_1_8, 8_3, 8_2_7_7_5, 4_7, 4_7_9, 9, 1_5_1_7, 7_3, 5_3_8_9_4, 3_3_3, 8_0_5_8_1, 1_1_0_1_1_7, 1_8_8_1_1, 5_2_5_6, 1_2_9_5, 5_1, 1_5_2_5_2_6, 2_9_7, 7_9_8_6, 3_9_0, 1_2_4_4_1_6, 5_3_8, 3_5_4_3_1, 2_1_4, 9_8, 1_5_0_4_4, 2_5_7_3_7, 1_3_6, 7_1_0_8, 4_3_7_0_1, 2_3, 7_5_6, 1_3_5_3_5_5, 7, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [2_5_0_0_0_4, 5_8_1, 6_3_7_7_3, 1_1_9_4_5_5, 6, 1_4_7_7_9_7, 8_8_2_0_3, 7, 6_4_5, 7_0, 2_1, 3_2_8_5, 1_0_2_6_9, 5, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], "attention_mask": [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
        # fmt: on

        self.tokenizer_integration_test_util(
            expected_encoding=lowerCAmelCase__		,				model_name="facebook/mbart-large-50"		,				revision="d3913889c59cd5c9e456b269c376325eabad57e2"		,				)







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




        if not self.test_slow_tokenizer:
          # as we don't have a slow version, we can't compare the outputs between slow and fast versions
          return

        a__     : Union[str, Any]  =(self.rust_tokenizer_class, "hf-internal-testing/tiny-random-mbart50", {})
        for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
          with self.subTest(F'''{tokenizer.__class__.__name__} ({pretrained_name})'''							):
            a__     : str  =self.rust_tokenizer_class.from_pretrained(lowerCAmelCase__		,				**lowerCAmelCase__							)
            a__     : Tuple  =self.tokenizer_class.from_pretrained(lowerCAmelCase__		,				**lowerCAmelCase__							)

            a__     : str  =tempfile.mkdtemp()

            a__     : str  =tokenizer_r.save_pretrained(lowerCAmelCase__							)
            a__     : List[str]  =tokenizer_p.save_pretrained(lowerCAmelCase__							)

            # Checks it save with the same files + the tokenizer.json file for the fast one
            self.assertTrue(any("tokenizer.json" in f for f in tokenizer_r_files							)							)
            a__     : List[str]  =tuple(f for f in tokenizer_r_files if "tokenizer.json" not in f							)
            self.assertSequenceEqual(lowerCAmelCase__		,				lowerCAmelCase__							)

            # Checks everything loads correctly in the same way
            a__     : Tuple  =tokenizer_r.from_pretrained(lowerCAmelCase__							)
            a__     : Tuple  =tokenizer_p.from_pretrained(lowerCAmelCase__							)

            # Check special tokens are set accordingly on Rust and Python
            for key in tokenizer_pp.special_tokens_map:
              self.assertTrue(hasattr(lowerCAmelCase__		,				lowerCAmelCase__							)							)
              # self.assertEqual(getattr(tokenizer_rp, key), getattr(tokenizer_pp, key))
              # self.assertEqual(getattr(tokenizer_rp, key + "_id"), getattr(tokenizer_pp, key + "_id"))

            shutil.rmtree(lowerCAmelCase__							)

            # Save tokenizer rust, legacy_format=True
            a__     : str  =tempfile.mkdtemp()

            a__     : str  =tokenizer_r.save_pretrained(lowerCAmelCase__		,				legacy_format=lowerCAmelCase__							)
            a__     : Dict  =tokenizer_p.save_pretrained(lowerCAmelCase__							)

            # Checks it save with the same files
            self.assertSequenceEqual(lowerCAmelCase__		,				lowerCAmelCase__							)

            # Checks everything loads correctly in the same way
            a__     : Tuple  =tokenizer_r.from_pretrained(lowerCAmelCase__							)
            a__     : str  =tokenizer_p.from_pretrained(lowerCAmelCase__							)

            # Check special tokens are set accordingly on Rust and Python
            for key in tokenizer_pp.special_tokens_map:
              self.assertTrue(hasattr(lowerCAmelCase__		,				lowerCAmelCase__							)							)

            shutil.rmtree(lowerCAmelCase__							)

            # Save tokenizer rust, legacy_format=False
            a__     : Optional[Any]  =tempfile.mkdtemp()

            a__     : List[Any]  =tokenizer_r.save_pretrained(lowerCAmelCase__		,				legacy_format=lowerCAmelCase__							)
            a__     : Optional[int]  =tokenizer_p.save_pretrained(lowerCAmelCase__							)

            # Checks it saved the tokenizer.json file
            self.assertTrue(any("tokenizer.json" in f for f in tokenizer_r_files							)							)

            # Checks everything loads correctly in the same way
            a__     : Tuple  =tokenizer_r.from_pretrained(lowerCAmelCase__							)
            a__     : Tuple  =tokenizer_p.from_pretrained(lowerCAmelCase__							)

            # Check special tokens are set accordingly on Rust and Python
            for key in tokenizer_pp.special_tokens_map:
              self.assertTrue(hasattr(lowerCAmelCase__		,				lowerCAmelCase__							)							)

            shutil.rmtree(lowerCAmelCase__							)







@require_torch
@require_sentencepiece
@require_tokenizers
class     __lowerCAmelCase	(   unittest.TestCase):
      _lowercase    : List[Any]            =      """facebook/mbart-large-50-one-to-many-mmt"""
      _lowercase    : str            =      [
          """ UN Chief Says There Is No Military Solution in Syria""",
          """ Secretary-General Ban Ki-moon says his response to Russia's stepped up military support for Syria is that \"there is no military solution\" to the nearly five-year conflict and more weapons will only worsen the violence and misery for millions of people.""",
      ]
      _lowercase    : List[str]            =      [
          """Şeful ONU declară că nu există o soluţie militară în Siria""",
          """Secretarul General Ban Ki-moon declară că răspunsul său la intensificarea sprijinului militar al Rusiei"""
          """ pentru Siria este că \"nu există o soluţie militară\" la conflictul de aproape cinci ani şi că noi arme nu vor"""
          """ face decât să înrăutăţească violenţele şi mizeria pentru milioane de oameni.""",
      ]
      _lowercase    : Dict            =      [EN_CODE, 8274, 12_7873, 2_5916, 7, 8622, 2071, 438, 6_7485, 53, 18_7895, 23, 5_1712, 2]







      @classmethod
      def        _lowercase  (   cls							)       ->						str:



        '''simple docstring'''




        a__     : MBartaaTokenizer  =MBartaaTokenizer.from_pretrained(
            cls.checkpoint_name		,				src_lang="en_XX"		,				tgt_lang="ro_RO"							)
        a__     : str  =1
        return cls







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ar_AR"]		,				2_5_0_0_0_1							)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["en_EN"]		,				2_5_0_0_0_4							)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["ro_RO"]		,				2_5_0_0_2_0							)
        self.assertEqual(self.tokenizer.fairseq_tokens_to_ids["mr_IN"]		,				2_5_0_0_3_8							)







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        a__     : Optional[int]  =self.tokenizer.batch_encode_plus(self.src_text							).input_ids[0]
        self.assertListEqual(self.expected_src_tokens		,				lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        self.assertIn(lowerCAmelCase__		,				self.tokenizer.all_special_ids							)
        a__     : List[Any]  =[RO_CODE, 8_8_4, 9_0_1_9, 9_6, 9, 9_1_6, 8_6_7_9_2, 3_6, 1_8_7_4_3, 1_5_5_9_6, 5, 2]
        a__     : Union[str, Any]  =self.tokenizer.decode(lowerCAmelCase__		,				skip_special_tokens=lowerCAmelCase__							)
        a__     : Dict  =self.tokenizer.decode(generated_ids[1:]		,				skip_special_tokens=lowerCAmelCase__							)
        self.assertEqual(lowerCAmelCase__		,				lowerCAmelCase__							)
        self.assertNotIn(self.tokenizer.eos_token		,				lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        a__     : Optional[Any]  =["this is gunna be a long sentence " * 2_0]
        assert isinstance(src_text[0]		,				lowerCAmelCase__							)
        a__     : Optional[Any]  =1_0
        a__     : List[Any]  =self.tokenizer(lowerCAmelCase__		,				max_length=lowerCAmelCase__		,				truncation=lowerCAmelCase__							).input_ids[0]
        self.assertEqual(ids[0]		,				lowerCAmelCase__							)
        self.assertEqual(ids[-1]		,				2							)
        self.assertEqual(len(lowerCAmelCase__							)		,				lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[Any]:



        '''simple docstring'''




        self.assertListEqual(self.tokenizer.convert_tokens_to_ids(["<mask>", "ar_AR"]							)		,				[2_5_0_0_5_3, 2_5_0_0_0_1]							)







      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Any  =tempfile.mkdtemp()
        a__     : List[Any]  =self.tokenizer.fairseq_tokens_to_ids
        self.tokenizer.save_pretrained(lowerCAmelCase__							)
        a__     : Any  =MBartaaTokenizer.from_pretrained(lowerCAmelCase__							)
        self.assertDictEqual(new_tok.fairseq_tokens_to_ids		,				lowerCAmelCase__							)







      @require_torch
      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




        a__     : List[str]  =self.tokenizer(self.src_text		,				text_target=self.tgt_text		,				padding=lowerCAmelCase__		,				return_tensors="pt"							)
        a__     : int  =shift_tokens_right(batch["labels"]		,				self.tokenizer.pad_token_id							)

        # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
        assert batch.input_ids[1][0] == EN_CODE
        assert batch.input_ids[1][-1] == 2
        assert batch.labels[1][0] == RO_CODE
        assert batch.labels[1][-1] == 2
        assert batch.decoder_input_ids[1][:2].tolist() == [2, RO_CODE]







      @require_torch
      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : int  =self.tokenizer(
            self.src_text		,				text_target=self.tgt_text		,				padding=lowerCAmelCase__		,				truncation=lowerCAmelCase__		,				max_length=len(self.expected_src_tokens							)		,				return_tensors="pt"		,				)
        a__     : Any  =shift_tokens_right(batch["labels"]		,				self.tokenizer.pad_token_id							)

        self.assertIsInstance(lowerCAmelCase__		,				lowerCAmelCase__							)

        self.assertEqual((2, 1_4)		,				batch.input_ids.shape							)
        self.assertEqual((2, 1_4)		,				batch.attention_mask.shape							)
        a__     : str  =batch.input_ids.tolist()[0]
        self.assertListEqual(self.expected_src_tokens		,				lowerCAmelCase__							)
        self.assertEqual(2		,				batch.decoder_input_ids[0, 0]							)  # decoder_start_token_id
        # Test that special tokens are reset
        self.assertEqual(self.tokenizer.prefix_tokens		,				[EN_CODE]							)
        self.assertEqual(self.tokenizer.suffix_tokens		,				[self.tokenizer.eos_token_id]							)







      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        a__     : Optional[Any]  =self.tokenizer(self.src_text		,				padding=lowerCAmelCase__		,				truncation=lowerCAmelCase__		,				max_length=3		,				return_tensors="pt"							)
        a__     : Any  =self.tokenizer(
            text_target=self.tgt_text		,				padding=lowerCAmelCase__		,				truncation=lowerCAmelCase__		,				max_length=1_0		,				return_tensors="pt"							)
        a__     : str  =targets["input_ids"]
        a__     : Union[str, Any]  =shift_tokens_right(lowerCAmelCase__		,				self.tokenizer.pad_token_id							)

        self.assertEqual(batch.input_ids.shape[1]		,				3							)
        self.assertEqual(batch.decoder_input_ids.shape[1]		,				1_0							)







      @require_torch
      def        _lowercase  (   self							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : str  =self.tokenizer._build_translation_inputs(
            "A test"		,				return_tensors="pt"		,				src_lang="en_XX"		,				tgt_lang="ar_AR"							)

        self.assertEqual(
            nested_simplify(lowerCAmelCase__							)		,				{
                # en_XX, A, test, EOS
                "input_ids": [[2_5_0_0_0_4, 6_2, 3_0_3_4, 2]],
                "attention_mask": [[1, 1, 1, 1]],
                # ar_AR
                "forced_bos_token_id": 2_5_0_0_0_1,
            }		,				)













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







      def __init__(   self		,				lowerCAmelCase__		,				lowerCAmelCase__=1_3		,				lowerCAmelCase__=7		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=True		,				lowerCAmelCase__=9_9		,				lowerCAmelCase__=3_2		,				lowerCAmelCase__=5		,				lowerCAmelCase__=4		,				lowerCAmelCase__=3_7		,				lowerCAmelCase__="gelu"		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=0.1		,				lowerCAmelCase__=5_1_2		,				lowerCAmelCase__=1_6		,				lowerCAmelCase__=2		,				lowerCAmelCase__=0.02		,				lowerCAmelCase__=3		,				lowerCAmelCase__=4		,				lowerCAmelCase__=None		,				)       ->						List[Any]:



        '''simple docstring'''




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







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        a__     : List[Any]  =ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size							)

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

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

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

        a__     : Tuple  =self.get_config()

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







      def        _lowercase  (   self							)       ->						Tuple:



        '''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        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[Any]:



        '''simple docstring'''




        a__     : Tuple  =NystromformerModel(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : Optional[Any]  =model(lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : str  =model(lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__							)
        a__     : Optional[int]  =model(lowerCAmelCase__							)
        self.parent.assertEqual(result.last_hidden_state.shape		,				(self.batch_size, self.seq_length, self.hidden_size)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						int:



        '''simple docstring'''




        a__     : int  =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        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Any:



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerForQuestionAnswering(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : 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        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



        '''simple docstring'''




        a__     : Optional[Any]  =self.num_labels
        a__     : Dict  =NystromformerForSequenceClassification(lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        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_labels)							)







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						List[str]:



        '''simple docstring'''




        a__     : Tuple  =self.num_labels
        a__     : List[str]  =NystromformerForTokenClassification(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[Any]  =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        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        a__     : List[Any]  =self.num_choices
        a__     : Optional[Any]  =NystromformerForMultipleChoice(config=lowerCAmelCase__							)
        model.to(lowerCAmelCase__							)
        model.eval()
        a__     : List[str]  =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__     : List[Any]  =input_mask.unsqueeze(1							).expand(-1		,				self.num_choices		,				-1							).contiguous()
        a__     : Dict  =model(
            lowerCAmelCase__		,				attention_mask=lowerCAmelCase__		,				token_type_ids=lowerCAmelCase__		,				labels=lowerCAmelCase__		,				)
        self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)							)







      def        _lowercase  (   self							)       ->						Tuple:



        '''simple docstring'''




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







@require_torch
class     __lowerCAmelCase	(   UpperCamelCase__						,    UpperCamelCase__						,    unittest.TestCase):
      _lowercase    : int            =      (
          (
              NystromformerModel,
              NystromformerForMaskedLM,
              NystromformerForMultipleChoice,
              NystromformerForQuestionAnswering,
              NystromformerForSequenceClassification,
              NystromformerForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      _lowercase    : Union[str, Any]            =      (
          {
              """feature-extraction""": NystromformerModel,
              """fill-mask""": NystromformerForMaskedLM,
              """question-answering""": NystromformerForQuestionAnswering,
              """text-classification""": NystromformerForSequenceClassification,
              """token-classification""": NystromformerForTokenClassification,
              """zero-shot""": NystromformerForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      _lowercase    : Union[str, Any]            =      False
      _lowercase    : Union[str, Any]            =      False







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        a__     : Optional[int]  =NystromformerModelTester(self							)
        a__     : Optional[int]  =ConfigTester(self		,				config_class=lowerCAmelCase__		,				hidden_size=3_7							)







      def        _lowercase  (   self							)       ->						int:



        '''simple docstring'''




        self.config_tester.run_common_tests()







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : List[str]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          a__     : int  =type
          self.model_tester.create_and_check_model(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Optional[Any]  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Dict:



        '''simple docstring'''




        a__     : str  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						List[str]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Union[str, Any]:



        '''simple docstring'''




        a__     : Tuple  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*lowerCAmelCase__							)







      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : int  =self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase__							)







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




        for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__     : int  =NystromformerModel.from_pretrained(lowerCAmelCase__							)
          self.assertIsNotNone(lowerCAmelCase__							)







@require_torch
class     __lowerCAmelCase	(   unittest.TestCase):







      @slow
      def        _lowercase  (   self							)       ->						str:



        '''simple docstring'''




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

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

        a__     : List[str]  =torch.Size((1, 6, 7_6_8)							)
        self.assertEqual(output.shape		,				lowerCAmelCase__							)

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







      @slow
      def        _lowercase  (   self							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Union[str, Any]  ="the [MASK] of Belgium is Brussels"

        a__     : str  =AutoTokenizer.from_pretrained("uw-madison/nystromformer-512"							)
        a__     : int  =NystromformerForMaskedLM.from_pretrained("uw-madison/nystromformer-512"							)

        a__     : List[Any]  =tokenizer(lowerCAmelCase__		,				return_tensors="pt"							)

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

        a__     : List[str]  =token_logits[:, 2, :].argmax(-1							)[0]

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






