test_tokenization_t5.py 8.17 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# coding=utf-8
# Copyright 2018 Google T5 Authors and HuggingFace Inc. 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.
Aymeric Augustin's avatar
Aymeric Augustin committed
15

thomwolf's avatar
thomwolf committed
16
17

import os
18
import unittest
thomwolf's avatar
thomwolf committed
19

20
from transformers import BatchEncoding
21
from transformers.file_utils import cached_property
22
from transformers.testing_utils import _torch_available
23
24
from transformers.tokenization_t5 import T5Tokenizer, T5TokenizerFast
from transformers.tokenization_xlnet import SPIECE_UNDERLINE
thomwolf's avatar
thomwolf committed
25

26
from .test_tokenization_common import TokenizerTesterMixin
thomwolf's avatar
thomwolf committed
27

Aymeric Augustin's avatar
Aymeric Augustin committed
28

29
30
SAMPLE_VOCAB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures/test_sentencepiece.model")

31
32
FRAMEWORK = "pt" if _torch_available else "tf"

thomwolf's avatar
thomwolf committed
33

34
class T5TokenizationTest(TokenizerTesterMixin, unittest.TestCase):
thomwolf's avatar
thomwolf committed
35
36

    tokenizer_class = T5Tokenizer
37
38
    rust_tokenizer_class = T5TokenizerFast
    test_rust_tokenizer = True
thomwolf's avatar
thomwolf committed
39
40

    def setUp(self):
Julien Chaumond's avatar
Julien Chaumond committed
41
        super().setUp()
thomwolf's avatar
thomwolf committed
42
43

        # We have a SentencePiece fixture for testing
44
        tokenizer = T5Tokenizer(SAMPLE_VOCAB)
thomwolf's avatar
thomwolf committed
45
46
47
        tokenizer.save_pretrained(self.tmpdirname)

    def test_full_tokenizer(self):
48
        tokenizer = T5Tokenizer(SAMPLE_VOCAB)
thomwolf's avatar
thomwolf committed
49

50
51
        tokens = tokenizer.tokenize("This is a test")
        self.assertListEqual(tokens, ["鈻乀his", "鈻乮s", "鈻乤", "鈻乼", "est"])
thomwolf's avatar
thomwolf committed
52

53
        self.assertListEqual(tokenizer.convert_tokens_to_ids(tokens), [285, 46, 10, 170, 382])
thomwolf's avatar
thomwolf committed
54

55
        tokens = tokenizer.tokenize("I was born in 92000, and this is fals茅.")
thomwolf's avatar
thomwolf committed
56
        self.assertListEqual(
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
            tokens,
            [
                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",
                "茅",
                ".",
            ],
        )
        ids = tokenizer.convert_tokens_to_ids(tokens)
        self.assertListEqual(ids, [8, 21, 84, 55, 24, 19, 7, 0, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 0, 4])
thomwolf's avatar
thomwolf committed
84
85

        back_tokens = tokenizer.convert_ids_to_tokens(ids)
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
        self.assertListEqual(
            back_tokens,
            [
                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>",
                ".",
            ],
        )
112

113
114
115
116
    @cached_property
    def t5_base_tokenizer(self):
        return T5Tokenizer.from_pretrained("t5-base")

117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
    @cached_property
    def t5_base_tokenizer_fast(self):
        return T5TokenizerFast.from_pretrained("t5-base")

    def get_tokenizer(self, **kwargs) -> T5Tokenizer:
        return self.tokenizer_class.from_pretrained(self.tmpdirname, pad_token=None, **kwargs)

    def get_rust_tokenizer(self, **kwargs) -> T5TokenizerFast:
        return self.rust_tokenizer_class.from_pretrained(self.tmpdirname, pad_token=None, **kwargs)

    def test_rust_and_python_full_tokenizers(self):
        if not self.test_rust_tokenizer:
            return

        tokenizer = self.get_tokenizer()
        rust_tokenizer = self.get_rust_tokenizer()

        sequence = "I was born in 92000, and this is fals茅."

        tokens = tokenizer.tokenize(sequence)
        rust_tokens = rust_tokenizer.tokenize(sequence)
        self.assertListEqual(tokens, rust_tokens)

        ids = tokenizer.encode(sequence, add_special_tokens=False)
        rust_ids = rust_tokenizer.encode(sequence, add_special_tokens=False)
        self.assertListEqual(ids, rust_ids)

        rust_tokenizer = self.get_rust_tokenizer()
        ids = tokenizer.encode(sequence)
        rust_ids = rust_tokenizer.encode(sequence)
        self.assertListEqual(ids, rust_ids)

149
150
151
152
153
154
    def test_eos_treatment(self):
        tokenizer = self.t5_base_tokenizer
        batch_with_eos_added = tokenizer(["hi</s>", "I went to the gym</s>", "</s>"])
        batch_without_eos_added = tokenizer(["hi", "I went to the gym", ""])
        self.assertListEqual(batch_with_eos_added["input_ids"], batch_without_eos_added["input_ids"])

155
    def test_prepare_seq2seq_batch(self):
156
        tokenizer = self.t5_base_tokenizer
157
        src_text = ["A long paragraph for summarization.", "Another paragraph for summarization."]
158
159
160
161
        tgt_text = [
            "Summary of the text.",
            "Another summary.",
        ]
162
        expected_src_tokens = [71, 307, 8986, 21, 4505, 1635, 1707, 5, tokenizer.eos_token_id]
Lysandre's avatar
Lysandre committed
163
164
165
166
167
        batch = tokenizer.prepare_seq2seq_batch(
            src_text,
            tgt_texts=tgt_text,
            return_tensors=FRAMEWORK,
        )
168
169
170
        self.assertIsInstance(batch, BatchEncoding)
        result = list(batch.input_ids.numpy()[0])
        self.assertListEqual(expected_src_tokens, result)
171

172
173
        self.assertEqual((2, 9), batch.input_ids.shape)
        self.assertEqual((2, 9), batch.attention_mask.shape)
174

175
    def test_empty_target_text(self):
176
        tokenizer = self.t5_base_tokenizer
177
        src_text = ["A long paragraph for summarization.", "Another paragraph for summarization."]
178
179
180
181
182
183
184
185
        batch = tokenizer.prepare_seq2seq_batch(src_text, return_tensors=FRAMEWORK)
        # check if input_ids are returned and no decoder_input_ids
        self.assertIn("input_ids", batch)
        self.assertIn("attention_mask", batch)
        self.assertNotIn("decoder_input_ids", batch)
        self.assertNotIn("decoder_attention_mask", batch)

    def test_max_target_length(self):
186
        tokenizer = self.t5_base_tokenizer
187
        src_text = ["A short paragraph for summarization.", "Another short paragraph for summarization."]
188
189
190
191
192
193
194
        tgt_text = [
            "Summary of the text.",
            "Another summary.",
        ]
        batch = tokenizer.prepare_seq2seq_batch(
            src_text, tgt_texts=tgt_text, max_target_length=32, padding="max_length", return_tensors=FRAMEWORK
        )
195
        self.assertEqual(32, batch["labels"].shape[1])
196
197
198
199
200

        # test None max_target_length
        batch = tokenizer.prepare_seq2seq_batch(
            src_text, tgt_texts=tgt_text, max_length=32, padding="max_length", return_tensors=FRAMEWORK
        )
201
        self.assertEqual(32, batch["labels"].shape[1])
202
203

    def test_outputs_not_longer_than_maxlen(self):
204
        tokenizer = self.t5_base_tokenizer
205
206
207
208
209
210
211
212

        batch = tokenizer.prepare_seq2seq_batch(
            ["I am a small frog" * 1000, "I am a small frog"], return_tensors=FRAMEWORK
        )
        self.assertIsInstance(batch, BatchEncoding)
        self.assertEqual(batch.input_ids.shape, (2, 512))

    def test_eos_in_input(self):
213
        tokenizer = self.t5_base_tokenizer
214
        src_text = ["A long paragraph for summarization. </s>"]
215
        tgt_text = ["Summary of the text. </s>"]
216
        expected_src_tokens = [71, 307, 8986, 21, 4505, 1635, 1707, 5, 1]
217
        expected_tgt_tokens = [20698, 13, 8, 1499, 5, 1]
218
219
220
221

        batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, return_tensors=FRAMEWORK)

        src_ids = list(batch.input_ids.numpy()[0])
222
        tgt_ids = list(batch.labels.numpy()[0])
223
224
225

        self.assertEqual(expected_src_tokens, src_ids)
        self.assertEqual(expected_tgt_tokens, tgt_ids)