This repository was archived by the owner on Dec 16, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2.2k
/
Copy pathpretrained_transformer_indexer_test.py
213 lines (191 loc) · 10.8 KB
/
pretrained_transformer_indexer_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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
84
85
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
112
113
114
115
116
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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
import pytest
from allennlp.common import cached_transformers
from allennlp.common.testing import AllenNlpTestCase
from allennlp.data import Vocabulary
from allennlp.data.token_indexers import PretrainedTransformerIndexer
from allennlp.data.tokenizers import PretrainedTransformerTokenizer
class TestPretrainedTransformerIndexer(AllenNlpTestCase):
def test_as_array_produces_token_sequence_bert_uncased(self):
tokenizer = cached_transformers.get_tokenizer("bert-base-uncased")
allennlp_tokenizer = PretrainedTransformerTokenizer("bert-base-uncased")
indexer = PretrainedTransformerIndexer(model_name="bert-base-uncased")
string_specials = "[CLS] AllenNLP is great [SEP]"
string_no_specials = "AllenNLP is great"
tokens = tokenizer.tokenize(string_specials)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
# tokens tokenized with our pretrained tokenizer have indices in them
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids
def test_as_array_produces_token_sequence_bert_cased(self):
tokenizer = cached_transformers.get_tokenizer("bert-base-cased")
allennlp_tokenizer = PretrainedTransformerTokenizer("bert-base-cased")
indexer = PretrainedTransformerIndexer(model_name="bert-base-cased")
string_specials = "[CLS] AllenNLP is great [SEP]"
string_no_specials = "AllenNLP is great"
tokens = tokenizer.tokenize(string_specials)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
# tokens tokenized with our pretrained tokenizer have indices in them
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids
def test_as_array_produces_token_sequence_bert_cased_sentence_pair(self):
tokenizer = cached_transformers.get_tokenizer("bert-base-cased")
allennlp_tokenizer = PretrainedTransformerTokenizer(
"bert-base-cased", add_special_tokens=False
)
indexer = PretrainedTransformerIndexer(model_name="bert-base-cased")
default_format = "[CLS] AllenNLP is great! [SEP] Really it is! [SEP]"
tokens = tokenizer.tokenize(default_format)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
allennlp_tokens = allennlp_tokenizer.add_special_tokens(
allennlp_tokenizer.tokenize("AllenNLP is great!"),
allennlp_tokenizer.tokenize("Really it is!"),
)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids
def test_as_array_produces_token_sequence_roberta(self):
tokenizer = cached_transformers.get_tokenizer("roberta-base")
allennlp_tokenizer = PretrainedTransformerTokenizer("roberta-base")
indexer = PretrainedTransformerIndexer(model_name="roberta-base")
string_specials = "<s>AllenNLP is great</s>"
string_no_specials = "AllenNLP is great"
tokens = tokenizer.tokenize(string_specials)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
# tokens tokenized with our pretrained tokenizer have indices in them
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids
def test_as_array_produces_token_sequence_roberta_sentence_pair(self):
tokenizer = cached_transformers.get_tokenizer("roberta-base")
allennlp_tokenizer = PretrainedTransformerTokenizer(
"roberta-base", add_special_tokens=False
)
indexer = PretrainedTransformerIndexer(model_name="roberta-base")
default_format = "<s>AllenNLP is great!</s></s>Really it is!</s>"
tokens = tokenizer.tokenize(default_format)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
allennlp_tokens = allennlp_tokenizer.add_special_tokens(
allennlp_tokenizer.tokenize("AllenNLP is great!"),
allennlp_tokenizer.tokenize("Really it is!"),
)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids, f"{allennlp_tokens}\n{tokens}"
@pytest.mark.parametrize("model_name", ["roberta-base", "bert-base-cased", "xlm-mlm-ende-1024"])
def test_transformers_vocab_sizes(self, model_name):
namespace = "tags"
tokenizer = cached_transformers.get_tokenizer(model_name)
allennlp_tokenizer = PretrainedTransformerTokenizer(model_name)
indexer = PretrainedTransformerIndexer(model_name=model_name, namespace=namespace)
allennlp_tokens = allennlp_tokenizer.tokenize("AllenNLP is great!")
vocab = Vocabulary()
# here we copy entire transformers vocab
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
del indexed
assert vocab.get_vocab_size(namespace=namespace) == tokenizer.vocab_size
def test_transformers_vocabs_added_correctly(self):
namespace, model_name = "tags", "roberta-base"
tokenizer = cached_transformers.get_tokenizer(model_name, use_fast=False)
allennlp_tokenizer = PretrainedTransformerTokenizer(model_name)
indexer = PretrainedTransformerIndexer(model_name=model_name, namespace=namespace)
allennlp_tokens = allennlp_tokenizer.tokenize("AllenNLP is great!")
vocab = Vocabulary()
# here we copy entire transformers vocab
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
del indexed
assert vocab.get_token_to_index_vocabulary(namespace=namespace) == tokenizer.encoder
def test_mask(self):
# We try these models, because
# - BERT pads tokens with 0
# - RoBERTa pads tokens with 1
# - GPT2 has no padding token, so we choose 0
for model in ["bert-base-uncased", "roberta-base", "gpt2"]:
allennlp_tokenizer = PretrainedTransformerTokenizer(model)
indexer = PretrainedTransformerIndexer(model_name=model)
string_no_specials = "AllenNLP is great"
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
expected_masks = [True] * len(indexed["token_ids"])
assert indexed["mask"] == expected_masks
max_length = 10
padding_lengths = {key: max_length for key in indexed.keys()}
padded_tokens = indexer.as_padded_tensor_dict(indexed, padding_lengths)
padding_length = max_length - len(indexed["mask"])
expected_masks = expected_masks + ([False] * padding_length)
assert len(padded_tokens["mask"]) == max_length
assert padded_tokens["mask"].tolist() == expected_masks
assert len(padded_tokens["token_ids"]) == max_length
pad_token_id = allennlp_tokenizer.tokenizer.pad_token_id
if pad_token_id is None:
pad_token_id = 0
padding_suffix = [pad_token_id] * padding_length
assert padded_tokens["token_ids"][-padding_length:].tolist() == padding_suffix
def test_long_sequence_splitting(self):
tokenizer = cached_transformers.get_tokenizer("bert-base-uncased")
allennlp_tokenizer = PretrainedTransformerTokenizer("bert-base-uncased")
indexer = PretrainedTransformerIndexer(model_name="bert-base-uncased", max_length=4)
string_specials = "[CLS] AllenNLP is great [SEP]"
string_no_specials = "AllenNLP is great"
tokens = tokenizer.tokenize(string_specials)
expected_ids = tokenizer.convert_tokens_to_ids(tokens)
assert len(expected_ids) == 7 # just to make sure it's what we're expecting
cls_id, sep_id = expected_ids[0], expected_ids[-1]
expected_ids = (
expected_ids[:3]
+ [sep_id, cls_id]
+ expected_ids[3:5]
+ [sep_id, cls_id]
+ expected_ids[5:]
)
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(allennlp_tokens, vocab)
assert indexed["token_ids"] == expected_ids
assert indexed["segment_concat_mask"] == [True] * len(expected_ids)
assert indexed["mask"] == [True] * 7 # original length
def test_type_ids_when_folding(self):
allennlp_tokenizer = PretrainedTransformerTokenizer(
"bert-base-uncased", add_special_tokens=False
)
indexer = PretrainedTransformerIndexer(model_name="bert-base-uncased", max_length=6)
first_string = "How do trees get online?"
second_string = "They log in!"
tokens = allennlp_tokenizer.add_special_tokens(
allennlp_tokenizer.tokenize(first_string), allennlp_tokenizer.tokenize(second_string)
)
vocab = Vocabulary()
indexed = indexer.tokens_to_indices(tokens, vocab)
assert min(indexed["type_ids"]) == 0
assert max(indexed["type_ids"]) == 1
@staticmethod
def _assert_tokens_equal(expected_tokens, actual_tokens):
for expected, actual in zip(expected_tokens, actual_tokens):
assert expected.text == actual.text
assert expected.text_id == actual.text_id
assert expected.type_id == actual.type_id
def test_indices_to_tokens(self):
allennlp_tokenizer = PretrainedTransformerTokenizer("bert-base-uncased")
indexer_max_length = PretrainedTransformerIndexer(
model_name="bert-base-uncased", max_length=4
)
indexer_no_max_length = PretrainedTransformerIndexer(model_name="bert-base-uncased")
string_no_specials = "AllenNLP is great"
allennlp_tokens = allennlp_tokenizer.tokenize(string_no_specials)
vocab = Vocabulary()
indexed = indexer_no_max_length.tokens_to_indices(allennlp_tokens, vocab)
tokens_from_indices = indexer_no_max_length.indices_to_tokens(indexed, vocab)
self._assert_tokens_equal(allennlp_tokens, tokens_from_indices)
indexed = indexer_max_length.tokens_to_indices(allennlp_tokens, vocab)
tokens_from_indices = indexer_max_length.indices_to_tokens(indexed, vocab)
# For now we are not removing special tokens introduced from max_length
sep_cls = [allennlp_tokens[-1], allennlp_tokens[0]]
expected = (
allennlp_tokens[:3] + sep_cls + allennlp_tokens[3:5] + sep_cls + allennlp_tokens[5:]
)
self._assert_tokens_equal(expected, tokens_from_indices)