test_seq2seq_lm.py 11.2 KB
Newer Older
1
2
3
4
5
import pytest
import torch

from copy import copy

6
7
from transformers import AutoTokenizer

8
9
10
11
from text_generation.pb import generate_pb2
from text_generation.models.seq2seq_lm import Seq2SeqLM, Seq2SeqLMBatch


12
13
14
15
16
17
18
19
20
21
22
23
24
25
@pytest.fixture(scope="session")
def mt0_small_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained(
        "bigscience/mt0-small", padding_side="left"
    )
    tokenizer.bos_token_id = 0
    return tokenizer


@pytest.fixture(scope="session")
def default_seq2seq_lm():
    return Seq2SeqLM("bigscience/mt0-small")


26
@pytest.fixture
27
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
28
29
30
31
32
    return generate_pb2.Request(
        id=0,
        inputs="Test",
        input_length=2,
        parameters=default_pb_parameters,
33
        stopping_parameters=default_pb_stop_parameters,
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
    )


@pytest.fixture
def default_pb_batch(default_pb_request):
    return generate_pb2.Batch(id=0, requests=[default_pb_request], size=1)


@pytest.fixture
def default_seq2seq_lm_batch(default_pb_batch, mt0_small_tokenizer):
    return Seq2SeqLMBatch.from_pb(
        default_pb_batch, mt0_small_tokenizer, torch.device("cpu")
    )


@pytest.fixture
def default_multi_requests_seq2seq_lm_batch(default_pb_request, mt0_small_tokenizer):
    req_0 = copy(default_pb_request)
    req_1 = default_pb_request
    req_1.id = 1
54
    req_1.stopping_parameters.max_new_tokens = 5
55
56
57
58
59
60
61

    batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
    return Seq2SeqLMBatch.from_pb(batch_pb, mt0_small_tokenizer, torch.device("cpu"))


def test_batch_from_pb(default_pb_batch, default_seq2seq_lm_batch):
    batch = default_seq2seq_lm_batch
62
    sequence_length = len(default_seq2seq_lm_batch.input_ids[0])
63
64
65
66

    assert batch.batch_id == default_pb_batch.id
    assert batch.requests == default_pb_batch.requests

67
    assert batch.input_ids.shape == (default_pb_batch.size, sequence_length)
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
    assert batch.input_ids[0][-2] == 4268
    assert batch.input_ids[0][-1] == 1
    assert torch.all(batch.input_ids[0][:-2] == 0)

    assert torch.all(batch.attention_mask[0][-2:] == 1)
    assert torch.all(batch.attention_mask[0][:-2] == 0)

    assert batch.decoder_input_ids.shape == (default_pb_batch.size, 1)
    assert batch.decoder_attention_mask is None
    assert batch.encoder_last_hidden_state is None

    assert batch.past_key_values is None

    assert batch.input_lengths == [2]
    assert batch.decoder_input_lengths == [1]

    assert batch.size == default_pb_batch.size
    assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == batch.size

    assert batch.max_input_length == batch.input_lengths[0]
    assert batch.max_decoder_input_length == batch.decoder_input_lengths[0]


def test_batch_concatenate_no_prefill(default_seq2seq_lm_batch):
    with pytest.raises(ValueError):
        Seq2SeqLMBatch.concatenate([default_seq2seq_lm_batch, default_seq2seq_lm_batch])


def test_seq2seq_lm_batch_type(default_seq2seq_lm):
    assert default_seq2seq_lm.batch_type == Seq2SeqLMBatch


def test_seq2seq_lm_generate_token(default_seq2seq_lm, default_seq2seq_lm_batch):
101
    sequence_length = len(default_seq2seq_lm_batch.input_ids[0])
102
    generated_texts, next_batch = default_seq2seq_lm.generate_token(
103
104
105
        default_seq2seq_lm_batch
    )

106
    assert generated_texts == []
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
    assert isinstance(next_batch, Seq2SeqLMBatch)

    assert torch.equal(next_batch.input_ids, default_seq2seq_lm_batch.input_ids)
    assert torch.equal(
        next_batch.attention_mask, default_seq2seq_lm_batch.attention_mask
    )
    assert next_batch.input_lengths == default_seq2seq_lm_batch.input_lengths
    assert next_batch.max_input_length == default_seq2seq_lm_batch.max_input_length
    assert (
        next_batch.next_token_choosers == default_seq2seq_lm_batch.next_token_choosers
    )
    assert next_batch.stopping_criterias == default_seq2seq_lm_batch.stopping_criterias

    assert next_batch.decoder_input_ids.shape == (next_batch.size, 2)
    assert next_batch.decoder_input_ids[0, 0] == 0
    assert next_batch.decoder_input_ids[0, 1] == 259
    assert next_batch.decoder_attention_mask is None
124
    assert next_batch.encoder_last_hidden_state.shape == (1, sequence_length, 512)
125
126
127
128
129
130
131
132
133
134
135
136

    assert next_batch.decoder_input_lengths == [2]
    assert next_batch.max_decoder_input_length == 2

    assert next_batch.past_key_values is not None
    assert all(
        [p[0].shape == (next_batch.size, 6, 1, 64) for p in next_batch.past_key_values]
    )
    assert all(
        [p[1].shape == (next_batch.size, 6, 1, 64) for p in next_batch.past_key_values]
    )
    assert all(
137
138
139
140
        [
            p[2].shape == (next_batch.size, 6, sequence_length, 64)
            for p in next_batch.past_key_values
        ]
141
142
    )
    assert all(
143
144
145
146
        [
            p[3].shape == (next_batch.size, 6, sequence_length, 64)
            for p in next_batch.past_key_values
        ]
147
148
149
150
151
152
153
154
    )


def test_seq2seq_lm_generate_token_completion(
    default_seq2seq_lm, default_seq2seq_lm_batch
):
    next_batch = default_seq2seq_lm_batch
    for _ in range(6):
155
156
        generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
        assert generated_texts == []
157

158
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
159
160
    assert next_batch is None

161
162
163
164
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few weeks"
    assert generated_texts[0].request == default_seq2seq_lm_batch.requests[0]
    assert generated_texts[0].generated_tokens == 7
165
166
167
168
169
170
171
172


def test_seq2seq_lm_generate_token_completion_multi(
    default_seq2seq_lm, default_multi_requests_seq2seq_lm_batch
):
    next_batch = default_multi_requests_seq2seq_lm_batch

    for i in range(4):
173
174
        generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
        assert generated_texts == []
175

176
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
177
178
    assert next_batch is not None

179
180
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few "
181
    assert (
182
183
        generated_texts[0].request
        == default_multi_requests_seq2seq_lm_batch.requests[1]
184
    )
185
    assert generated_texts[0].generated_tokens == 5
186

187
188
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
    assert generated_texts == []
189

190
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
191
192
    assert next_batch is None

193
194
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few weeks"
195
    assert (
196
197
        generated_texts[0].request
        == default_multi_requests_seq2seq_lm_batch.requests[0]
198
    )
199
    assert generated_texts[0].generated_tokens == 7
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293


def test_batch_concatenate(
    default_seq2seq_lm,
    default_seq2seq_lm_batch,
    default_multi_requests_seq2seq_lm_batch,
):
    next_batch_0 = default_seq2seq_lm_batch
    _, next_batch_0 = default_seq2seq_lm.generate_token(next_batch_0)
    _, next_batch_0 = default_seq2seq_lm.generate_token(next_batch_0)

    next_batch_1 = default_multi_requests_seq2seq_lm_batch
    _, next_batch_1 = default_seq2seq_lm.generate_token(next_batch_1)

    next_batch = Seq2SeqLMBatch.concatenate([next_batch_0, next_batch_1])

    assert next_batch.batch_id == 0

    assert torch.all(next_batch.input_ids[:, 0] == 4268)
    assert torch.all(next_batch.input_ids[:, 1] == 1)

    assert torch.all(next_batch.attention_mask == 1)

    assert torch.equal(
        next_batch.decoder_input_ids[0], next_batch_0.decoder_input_ids[0]
    )
    assert torch.all(next_batch.decoder_input_ids[1:, 0] == 0)
    assert torch.equal(
        next_batch.decoder_input_ids[1:, -2:], next_batch_1.decoder_input_ids
    )

    assert torch.all(next_batch.decoder_attention_mask[0] == 1)
    assert torch.all(next_batch.decoder_attention_mask[1:, 0] == 0)
    assert torch.all(next_batch.decoder_attention_mask[1:, -2:] == 1)

    assert torch.equal(
        next_batch.encoder_last_hidden_state[0],
        next_batch_0.encoder_last_hidden_state[0, -2:],
    )
    assert torch.equal(
        next_batch.encoder_last_hidden_state[1:],
        next_batch_1.encoder_last_hidden_state[:, -2:],
    )

    assert next_batch.input_lengths == [2, 2, 2]
    assert next_batch.decoder_input_lengths == [3, 2, 2]
    assert next_batch.max_input_length == 2
    assert next_batch.max_decoder_input_length == 3

    assert next_batch.requests[0] == next_batch_0.requests[0]
    assert next_batch.requests[1:] == next_batch_1.requests

    assert next_batch.next_token_choosers[0] == next_batch_0.next_token_choosers[0]
    assert next_batch.next_token_choosers[1:] == next_batch_1.next_token_choosers

    assert next_batch.stopping_criterias[0] == next_batch_0.stopping_criterias[0]
    assert next_batch.stopping_criterias[1:] == next_batch_1.stopping_criterias

    assert next_batch.past_key_values is not None
    assert all(
        [p[0].shape == (next_batch.size, 6, 2, 64) for p in next_batch.past_key_values]
    )
    assert all(
        [p[1].shape == (next_batch.size, 6, 2, 64) for p in next_batch.past_key_values]
    )
    assert all(
        [p[2].shape == (next_batch.size, 6, 2, 64) for p in next_batch.past_key_values]
    )
    assert all(
        [p[3].shape == (next_batch.size, 6, 2, 64) for p in next_batch.past_key_values]
    )

    for i, past in enumerate(next_batch.past_key_values):
        assert torch.equal(next_batch_0.past_key_values[i][0][0, :, -2:, :], past[0][0])
        assert torch.equal(
            next_batch_1.past_key_values[i][0][:, :, -1:, :], past[0][1:, :, -1:, :]
        )

        assert torch.equal(next_batch_0.past_key_values[i][1][0, :, -2:, :], past[1][0])
        assert torch.equal(
            next_batch_1.past_key_values[i][1][:, :, -1:, :], past[1][1:, :, -1:, :]
        )

        assert torch.equal(next_batch_0.past_key_values[i][2][0, :, -2:, :], past[2][0])
        assert torch.equal(
            next_batch_1.past_key_values[i][2][:, :, -2:, :], past[2][1:]
        )

        assert torch.equal(next_batch_0.past_key_values[i][3][0, :, -2:, :], past[3][0])
        assert torch.equal(
            next_batch_1.past_key_values[i][3][:, :, -2:, :], past[3][1:]
        )

    for _ in range(3):
294
295
        generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
        assert generated_texts == []
296

297
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
298
299
    assert next_batch is not None

300
301
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few "
302
    assert (
303
304
        generated_texts[0].request
        == default_multi_requests_seq2seq_lm_batch.requests[1]
305
    )
306
    assert generated_texts[0].generated_tokens == 5
307

308
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
309
310
    assert next_batch is not None

311
312
313
314
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few weeks"
    assert generated_texts[0].request == default_seq2seq_lm_batch.requests[0]
    assert generated_texts[0].generated_tokens == 7
315

316
    generated_texts, next_batch = default_seq2seq_lm.generate_token(next_batch)
317
318
    assert next_batch is None

319
320
    assert len(generated_texts) == 1
    assert generated_texts[0].output_text == "a few weeks"
321
    assert (
322
323
        generated_texts[0].request
        == default_multi_requests_seq2seq_lm_batch.requests[0]
324
    )
325
    assert generated_texts[0].generated_tokens == 7