2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license we have a sentencepiece fixture for testing test converttokentoid and convertidtotoken token s tokenid 1 self assertequalself gettokenizer converttokenstoidstoken tokenid self assertequalself gettokenizer convertidstotokenstokenid token def testgetvocabself vocabkeys listself gettokenizer getvocab keys self assertequalvocabkeys0 pad self assertequalvocabkeys1 s self assertequalvocabkeys104 unk102 self assertequallenvocabkeys 1103 def testvocabsizeself self assertequalself gettokenizer vocabsize 1103 def testmasktokensrustpegasusself rusttokenizer self rusttokenizerclass frompretrainedself tmpdirname pytokenizer self tokenizerclass frompretrainedself tmpdirname rawinputstr let s see which unk is the better unktoken11 one mask1 it seems like this mask2 was important s pad pad pad rustids rusttokenizerrawinputstr returntensorsnone addspecialtokensfalse inputids0 pyids pytokenizerrawinputstr returntensorsnone addspecialtokensfalse inputids0 self assertlistequalpyids rustids def testlargemasktokensself tokenizer self largetokenizer mask1 masks whole sentence while mask2 masks single word rawinputstr mask1 to ensure a mask2 flow of bank resolutions desiredresult 2 413 615 114 3 1971 113 1679 10710 107 1 ids tokenizerrawinputstr returntensorsnone inputids0 self assertlistequaldesiredresult ids def testlargetokenizersettingsself tokenizer self largetokenizer the tracebacks for the following asserts are better without messages or self assertequal assert tokenizer vocabsize 96103 assert tokenizer padtokenid 0 assert tokenizer eostokenid 1 assert tokenizer offset 103 assert tokenizer unktokenid tokenizer offset 2 105 assert tokenizer unktoken unk assert tokenizer modelmaxlength 1024 rawinputstr to ensure a smooth flow of bank resolutions desiredresult 413 615 114 2291 1971 113 1679 10710 107 1 ids tokenizerrawinputstr returntensorsnone inputids0 self assertlistequaldesiredresult ids assert tokenizer convertidstotokens0 1 2 3 pad s mask1 mask2 requiretorch def testlargeseq2seqtruncationself srctexts this is going to be way too long 150 short example tgttexts not super long but more than 5 tokens tiny batch self largetokenizersrctexts paddingtrue truncationtrue returntensorspt targets self largetokenizer texttargettgttexts maxlength5 paddingtrue truncationtrue returntensorspt assert batch inputids shape 2 1024 assert batch attentionmask shape 2 1024 assert targetsinputids shape 2 5 assert lenbatch 2 inputids attentionmask slow def testtokenizerintegrationself expectedencoding inputids 38979 143 18485 606 130 26669 87686 121 54189 1129 111 26669 87686 121 9114 14787 121 13249 158 592 956 121 14621 31576 143 62613 108 9688 930 43430 11562 62613 304 108 11443 897 108 9314 17415 63399 108 11443 7614 18316 118 4284 7148 12430 143 1400 25703 158 111 4284 7148 11772 143 21297 1064 158 122 204 3506 1754 1133 14787 1581 115 33224 4482 111 1355 110 29173 317 50833 108 20147 94665 111 77198 107 1 110 62613 117 638 112 1133 121 20098 1355 79050 13872 135 1596 53541 1352 141 13039 5542 124 302 518 111 268 2956 115 149 4427 107 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 139 1235 2799 18289 17780 204 109 9474 1296 107 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 attentionmask 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 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 fmt skip self tokenizerintegrationtestutil expectedencodingexpectedencoding modelnamegooglebigbirdpegasuslargearxiv revisionba85d0851d708441f91440d509690f1ab6353415 unittest skipwe have to use fromslow def testaddedtokensserializationself pass requiresentencepiece requiretokenizers class bigbirdpegasustokenizationtesttokenizertestermixin unittest testcase tokenizerclass pegasustokenizer rusttokenizerclass pegasustokenizerfast testrusttokenizer true testsentencepiece true def setupself super setup we have a sentencepiece fixture for testing tokenizer pegasustokenizersamplevocab offset0 masktokensentnone masktokenmask tokenizer savepretrainedself tmpdirname cachedproperty def largetokenizerself return pegasustokenizer frompretrainedgooglebigbirdpegasuslargearxiv def gettokenizerself kwargs pegasustokenizer return pegasustokenizer frompretrainedself tmpdirname kwargs def getinputoutputtextsself tokenizer return this is a test this is a test def testmasktokensrustpegasusself rusttokenizer self rusttokenizerclass frompretrainedself tmpdirname pytokenizer self tokenizerclass frompretrainedself tmpdirname rawinputstr let s see which unk is the better unktoken one mask it seems like this mask was important s pad pad pad rustids rusttokenizerrawinputstr returntensorsnone addspecialtokensfalse inputids0 pyids pytokenizerrawinputstr returntensorsnone addspecialtokensfalse inputids0 self assertlistequalpyids rustids requiretorch def testlargeseq2seqtruncationself srctexts this is going to be way too long 1000 short example tgttexts not super long but more than 5 tokens tiny batch self largetokenizersrctexts paddingtrue truncationtrue returntensorspt targets self largetokenizer texttargettgttexts maxlength5 paddingtrue truncationtrue returntensorspt assert batch inputids shape 2 4096 assert batch attentionmask shape 2 4096 assert targetsinputids shape 2 5 assert lenbatch 2 inputids attentionmask def testequivalencetoorigtokenizerself teststr this is an example string that is used to test the original tf implementation against the hf implementation tokenids self largetokenizerteststr inputids self assertlistequal tokenids 182 117 142 587 4211 120 117 263 112 804 109 856 25016 3137 464 109 26955 3137 1 unittest skipwe have to use fromslow def testaddedtokensserializationself pass 2020 the huggingface team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license we have a sentencepiece fixture for testing test _convert_token_to_id and _convert_id_to_token mask_1 masks whole sentence while mask_2 masks single word the tracebacks for the following asserts are better without messages or self assertequal input_ids attention_mask fmt skip unittest skip we have to use from_slow def test_added_tokens_serialization self pass we have a sentencepiece fixture for testing input_ids attention_mask to run with original tf tokenizer wget https github com google research bigbird raw master bigbird vocab pegasus model pip install tensorflow text import tensorflow compat v2 as tf import tensorflow_text as tft vocab_file pegasus model tf enable_v2_behavior test_str this is an example string that is used to test the original tf implementation against the hf implementation tokenizer tft sentencepiecetokenizer model tf io gfile gfile vocab_file rb read tokenizer tokenize test_str unittest skip we have to use from_slow def test_added_tokens_serialization self pass
import unittest

from transformers import PegasusTokenizer, PegasusTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, require_torch, slow
from transformers.utils import cached_property

from ...test_tokenization_common import TokenizerTesterMixin


SAMPLE_VOCAB = get_tests_dir("fixtures/test_sentencepiece_no_bos.model")


@require_sentencepiece
@require_tokenizers
class PegasusTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
    tokenizer_class = PegasusTokenizer
    rust_tokenizer_class = PegasusTokenizerFast
    test_rust_tokenizer = True
    test_sentencepiece = True

    def setUp(self):
        super().setUp()

        
        tokenizer = PegasusTokenizer(SAMPLE_VOCAB)
        tokenizer.save_pretrained(self.tmpdirname)

    @cached_property
    def _large_tokenizer(self):
        return PegasusTokenizer.from_pretrained("google/pegasus-large")

    def get_tokenizer(self, **kwargs) -> PegasusTokenizer:
        return PegasusTokenizer.from_pretrained(self.tmpdirname, **kwargs)

    def get_input_output_texts(self, tokenizer):
        return ("This is a test", "This is a test")

    def test_convert_token_and_id(self):
        
        token = "</s>"
        token_id = 1

        self.assertEqual(self.get_tokenizer().convert_tokens_to_ids(token), token_id)
        self.assertEqual(self.get_tokenizer().convert_ids_to_tokens(token_id), token)

    def test_get_vocab(self):
        vocab_keys = list(self.get_tokenizer().get_vocab().keys())

        self.assertEqual(vocab_keys[0], "<pad>")
        self.assertEqual(vocab_keys[1], "</s>")
        self.assertEqual(vocab_keys[104], "<unk_102>")
        self.assertEqual(len(vocab_keys), 1_103)

    def test_vocab_size(self):
        self.assertEqual(self.get_tokenizer().vocab_size, 1_103)

    def test_mask_tokens_rust_pegasus(self):
        rust_tokenizer = self.rust_tokenizer_class.from_pretrained(self.tmpdirname)
        py_tokenizer = self.tokenizer_class.from_pretrained(self.tmpdirname)
        raw_input_str = (
            "Let's see which <unk> is the better <unk_token_11> one <mask_1> It seems like this <mask_2> was important"
            " </s> <pad> <pad> <pad>"
        )
        rust_ids = rust_tokenizer([raw_input_str], return_tensors=None, add_special_tokens=False).input_ids[0]
        py_ids = py_tokenizer([raw_input_str], return_tensors=None, add_special_tokens=False).input_ids[0]
        self.assertListEqual(py_ids, rust_ids)

    def test_large_mask_tokens(self):
        tokenizer = self._large_tokenizer
        
        raw_input_str = "<mask_1> To ensure a <mask_2> flow of bank resolutions."
        desired_result = [2, 413, 615, 114, 3, 1971, 113, 1679, 10710, 107, 1]
        ids = tokenizer([raw_input_str], return_tensors=None).input_ids[0]
        self.assertListEqual(desired_result, ids)

    def test_large_tokenizer_settings(self):
        tokenizer = self._large_tokenizer
        
        assert tokenizer.vocab_size == 96103
        assert tokenizer.pad_token_id == 0
        assert tokenizer.eos_token_id == 1
        assert tokenizer.offset == 103
        assert tokenizer.unk_token_id == tokenizer.offset + 2 == 105
        assert tokenizer.unk_token == "<unk>"
        assert tokenizer.model_max_length == 1024
        raw_input_str = "To ensure a smooth flow of bank resolutions."
        desired_result = [413, 615, 114, 2291, 1971, 113, 1679, 10710, 107, 1]
        ids = tokenizer([raw_input_str], return_tensors=None).input_ids[0]
        self.assertListEqual(desired_result, ids)
        assert tokenizer.convert_ids_to_tokens([0, 1, 2, 3]) == ["<pad>", "</s>", "<mask_1>", "<mask_2>"]

    @require_torch
    def test_large_seq2seq_truncation(self):
        src_texts = ["This is going to be way too long." * 150, "short example"]
        tgt_texts = ["not super long but more than 5 tokens", "tiny"]
        batch = self._large_tokenizer(src_texts, padding=True, truncation=True, return_tensors="pt")
        targets = self._large_tokenizer(
            text_target=tgt_texts, max_length=5, padding=True, truncation=True, return_tensors="pt"
        )

        assert batch.input_ids.shape == (2, 1024)
        assert batch.attention_mask.shape == (2, 1024)
        assert targets["input_ids"].shape == (2, 5)
        assert len(batch) == 2  

    @slow
    def test_tokenizer_integration(self):
        expected_encoding = {'input_ids': [[38979, 143, 18485, 606, 130, 26669, 87686, 121, 54189, 1129, 111, 26669, 87686, 121, 9114, 14787, 121, 13249, 158, 592, 956, 121, 14621, 31576, 143, 62613, 108, 9688, 930, 43430, 11562, 62613, 304, 108, 11443, 897, 108, 9314, 17415, 63399, 108, 11443, 7614, 18316, 118, 4284, 7148, 12430, 143, 1400, 25703, 158, 111, 4284, 7148, 11772, 143, 21297, 1064, 158, 122, 204, 3506, 1754, 1133, 14787, 1581, 115, 33224, 4482, 111, 1355, 110, 29173, 317, 50833, 108, 20147, 94665, 111, 77198, 107, 1], [110, 62613, 117, 638, 112, 1133, 121, 20098, 1355, 79050, 13872, 135, 1596, 53541, 1352, 141, 13039, 5542, 124, 302, 518, 111, 268, 2956, 115, 149, 4427, 107, 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], [139, 1235, 2799, 18289, 17780, 204, 109, 9474, 1296, 107, 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]], '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, 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, 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]]}  

        self.tokenizer_integration_test_util(
            expected_encoding=expected_encoding,
            model_name="google/bigbird-pegasus-large-arxiv",
            revision="ba85d0851d708441f91440d509690f1ab6353415",
        )

    
    
    


@require_sentencepiece
@require_tokenizers
class BigBirdPegasusTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
    tokenizer_class = PegasusTokenizer
    rust_tokenizer_class = PegasusTokenizerFast
    test_rust_tokenizer = True
    test_sentencepiece = True

    def setUp(self):
        super().setUp()

        
        tokenizer = PegasusTokenizer(SAMPLE_VOCAB, offset=0, mask_token_sent=None, mask_token="[MASK]")
        tokenizer.save_pretrained(self.tmpdirname)

    @cached_property
    def _large_tokenizer(self):
        return PegasusTokenizer.from_pretrained("google/bigbird-pegasus-large-arxiv")

    def get_tokenizer(self, **kwargs) -> PegasusTokenizer:
        return PegasusTokenizer.from_pretrained(self.tmpdirname, **kwargs)

    def get_input_output_texts(self, tokenizer):
        return ("This is a test", "This is a test")

    def test_mask_tokens_rust_pegasus(self):
        rust_tokenizer = self.rust_tokenizer_class.from_pretrained(self.tmpdirname)
        py_tokenizer = self.tokenizer_class.from_pretrained(self.tmpdirname)
        raw_input_str = (
            "Let's see which <unk> is the better <unk_token> one [MASK] It seems like this [MASK] was important </s>"
            " <pad> <pad> <pad>"
        )
        rust_ids = rust_tokenizer([raw_input_str], return_tensors=None, add_special_tokens=False).input_ids[0]
        py_ids = py_tokenizer([raw_input_str], return_tensors=None, add_special_tokens=False).input_ids[0]
        self.assertListEqual(py_ids, rust_ids)

    @require_torch
    def test_large_seq2seq_truncation(self):
        src_texts = ["This is going to be way too long." * 1000, "short example"]
        tgt_texts = ["not super long but more than 5 tokens", "tiny"]
        batch = self._large_tokenizer(src_texts, padding=True, truncation=True, return_tensors="pt")
        targets = self._large_tokenizer(
            text_target=tgt_texts, max_length=5, padding=True, truncation=True, return_tensors="pt"
        )

        assert batch.input_ids.shape == (2, 4096)
        assert batch.attention_mask.shape == (2, 4096)
        assert targets["input_ids"].shape == (2, 5)
        assert len(batch) == 2  

    def test_equivalence_to_orig_tokenizer(self):
        

        test_str = (
            "This is an example string that is used to test the original TF implementation against the HF"
            " implementation"
        )

        token_ids = self._large_tokenizer(test_str).input_ids

        self.assertListEqual(
            token_ids,
            [182, 117, 142, 587, 4211, 120, 117, 263, 112, 804, 109, 856, 25016, 3137, 464, 109, 26955, 3137, 1],
        )