test_bloom.py 11.3 KB
Newer Older
1
2
3
4
import pytest
import torch

from copy import copy
5
from transformers import AutoTokenizer
6

7
8
9
from text_generation_server.pb import generate_pb2
from text_generation_server.models.causal_lm import CausalLMBatch
from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOM
10
11


12
13
14
15
16
17
18
19
20
21
@pytest.fixture(scope="session")
def default_bloom():
    return BLOOM("bigscience/bloom-560m")


@pytest.fixture(scope="session")
def bloom_560m_tokenizer():
    return AutoTokenizer.from_pretrained("bigscience/bloom-560m", padding_side="left")


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


@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_bloom_batch(default_pb_batch, bloom_560m_tokenizer):
    return BloomCausalLMBatch.from_pb(
        default_pb_batch, bloom_560m_tokenizer, torch.device("cpu")
    )


@pytest.fixture
def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer):
    req_0 = copy(default_pb_request)
    req_1 = default_pb_request
    req_1.id = 1
49
    req_1.stopping_parameters.max_new_tokens = 5
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

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


def test_batch_from_pb(default_pb_batch, default_bloom_batch):
    batch = default_bloom_batch

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

    assert len(batch.input_ids) == default_pb_batch.size
    assert batch.input_ids[0][-1] == 10264
    assert torch.all(batch.input_ids[0][:-1] == 3)

67
68
    assert batch.attention_mask[0][0] == 1
    assert torch.all(batch.attention_mask[0][1:] == 0)
69
70
71
72
73
74
75
76
77
78

    assert batch.past_key_values is None

    assert torch.equal(batch.input_ids, batch.all_input_ids[:, :, 0])

    assert batch.input_lengths == [1]

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

79
    assert batch.max_input_length == batch.input_lengths[0]
80
81
82
83
84
85
86
87
88
89
90
91


def test_batch_concatenate_no_prefill(default_bloom_batch):
    with pytest.raises(ValueError):
        BloomCausalLMBatch.concatenate([default_bloom_batch, default_bloom_batch])


def test_causal_lm_batch_type(default_bloom):
    assert default_bloom.batch_type == BloomCausalLMBatch


def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
92
    sequence_length = len(default_bloom_batch.all_input_ids[0])
93
    generations, next_batch = default_bloom.generate_token(default_bloom_batch)
94

95
    assert len(generations) == len(default_bloom_batch)
96
97
98
99
    assert isinstance(next_batch, CausalLMBatch)
    assert not next_batch.keys_head_dim_last

    assert len(next_batch.all_input_ids) == next_batch.size
100
101
    assert len(next_batch.all_input_ids[0]) == sequence_length + 1
    assert len(next_batch.attention_mask[0]) == 11
102
103
104
    assert torch.all(next_batch.all_input_ids[0][-2:] == 10264)
    assert torch.all(next_batch.all_input_ids[0][:-2] == 3)

105
106
    assert torch.all(next_batch.attention_mask[0][:2] == 1)
    assert torch.all(next_batch.attention_mask[0][2:] == 0)
107
108
109
110
111

    assert next_batch.input_ids.shape == (next_batch.size, 1)
    assert next_batch.input_ids[0, 0] == 10264

    assert next_batch.input_lengths == [2]
112
    assert next_batch.max_input_length == next_batch.input_lengths[0]
113
114

    assert next_batch.past_key_values is not None
115
116
117
118
119
120
    assert all(
        [p[0].shape == (16, 64, sequence_length) for p in next_batch.past_key_values]
    )
    assert all(
        [p[1].shape == (16, sequence_length, 64) for p in next_batch.past_key_values]
    )
121
122
123
124
125
    assert all([generation.generated_text is None for generation in generations])
    assert all([len(generation.prefill_tokens) == 1 for generation in generations])
    assert all([generation.token_id.item() == 10264 for generation in generations])
    assert all([generation.token_text == "Test" for generation in generations])
    assert generations[0].request_id == 0
126
127
128
129
130


def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch):
    next_batch = default_bloom_batch
    for _ in range(default_bloom_batch.stopping_criterias[0].max_new_tokens - 1):
131
132
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(default_bloom_batch)
133

134
    generations, next_batch = default_bloom.generate_token(next_batch)
135
136
    assert next_batch is None

137
    assert len(generations) == 1
OlivierDehaene's avatar
OlivierDehaene committed
138
    assert (
139
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
OlivierDehaene's avatar
OlivierDehaene committed
140
    )
141
    assert generations[0].request_id == default_bloom_batch.requests[0].id
142
    assert (
143
        generations[0].generated_text.generated_tokens
144
145
146
147
148
149
150
151
152
153
154
155
        == default_bloom_batch.stopping_criterias[0].max_new_tokens
    )


def test_causal_lm_generate_token_completion_multi(
    default_bloom, default_multi_requests_bloom_batch
):
    next_batch = default_multi_requests_bloom_batch

    for i in range(
        default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens - 1
    ):
156
157
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(default_multi_requests_bloom_batch)
158

159
    generations, next_batch = default_bloom.generate_token(next_batch)
160
161
    assert next_batch is not None

162
    assert len(generations) == 2
163
    assert generations[1].generated_text.text == "TestTestTestTestTest"
164
    assert (
165
166
167
168
        generations[1].request_id == default_multi_requests_bloom_batch.requests[1].id
    )
    assert (
        generations[1].generated_text.generated_tokens
169
170
171
172
173
174
175
176
        == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
    )

    for _ in range(
        default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
        - default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
        - 1
    ):
177
178
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(next_batch)
179

180
    generations, next_batch = default_bloom.generate_token(next_batch)
181
182
    assert next_batch is None

183
184
    assert len(generations) == 1
    assert (
185
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
186
    )
OlivierDehaene's avatar
OlivierDehaene committed
187
    assert (
188
        generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
OlivierDehaene's avatar
OlivierDehaene committed
189
    )
190
    assert (
191
        generations[0].generated_text.generated_tokens
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
        == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
    )


def test_batch_concatenate(
    default_bloom, default_bloom_batch, default_multi_requests_bloom_batch
):
    next_batch_0 = default_bloom_batch
    _, next_batch_0 = default_bloom.generate_token(next_batch_0)
    _, next_batch_0 = default_bloom.generate_token(next_batch_0)

    next_batch_1 = default_multi_requests_bloom_batch
    _, next_batch_1 = default_bloom.generate_token(next_batch_1)

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

    assert torch.equal(next_batch.all_input_ids[0], next_batch_0.all_input_ids[0])
    assert torch.equal(next_batch.all_input_ids[1], next_batch_1.all_input_ids[0])
    assert torch.equal(next_batch.all_input_ids[2], next_batch_1.all_input_ids[1])

212
213
214
215
216
217
218
    assert torch.all(
        next_batch.attention_mask[0, : -next_batch.padding_right_offset] == 1
    )
    assert torch.all(
        next_batch.attention_mask[1:, 1 : -next_batch.padding_right_offset] == 1
    )
    assert torch.all(next_batch.attention_mask[1:, 3:] == 0)
219
220
221
222
223

    assert next_batch.batch_id == 0
    assert torch.all(next_batch.input_ids == 10264)

    assert next_batch.input_lengths == [3, 2, 2]
224
    assert next_batch.max_input_length == 3
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

    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 == (3, 16, 64, 2) for p in next_batch.past_key_values])
    assert all([p[1].shape == (3, 16, 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][:, :, -2:], past[0][0])
        assert torch.equal(
            next_batch_1.past_key_values[i][0][:, :, -1:],
            past[0][1:, :, :, -1].reshape(-1, 64, 1),
        )

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

    for _ in range(
        default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens - 2
    ):
255
256
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(next_batch)
257

258
    generations, next_batch = default_bloom.generate_token(next_batch)
259
260
    assert next_batch is not None

261
    assert len(generations) == 3
262
    assert generations[2].generated_text.text == "TestTestTestTestTest"
263
    assert (
264
265
266
267
        generations[2].request_id == default_multi_requests_bloom_batch.requests[1].id
    )
    assert (
        generations[2].generated_text.generated_tokens
268
269
270
271
272
273
274
275
        == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
    )

    for _ in range(
        default_bloom_batch.stopping_criterias[0].max_new_tokens
        - default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
        - 2
    ):
276
277
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(next_batch)
278

279
    generations, next_batch = default_bloom.generate_token(next_batch)
280
281
    assert next_batch is not None

282
    assert len(generations) == 2
OlivierDehaene's avatar
OlivierDehaene committed
283
    assert (
284
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
OlivierDehaene's avatar
OlivierDehaene committed
285
    )
286
    assert generations[0].request_id == default_bloom_batch.requests[0].id
287
    assert (
288
        generations[0].generated_text.generated_tokens
289
290
291
292
293
294
295
296
297
        == default_bloom_batch.stopping_criterias[0].max_new_tokens
    )

    for _ in range(
        default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
        - default_bloom_batch.stopping_criterias[0].max_new_tokens
        - default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
        - 4
    ):
298
299
        generations, next_batch = default_bloom.generate_token(next_batch)
        assert len(generations) == len(next_batch)
300

301
    generations, next_batch = default_bloom.generate_token(next_batch)
302
303
    assert next_batch is None

304
305
    assert len(generations) == 1
    assert (
306
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
307
    )
OlivierDehaene's avatar
OlivierDehaene committed
308
    assert (
309
        generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
OlivierDehaene's avatar
OlivierDehaene committed
310
    )
311
    assert (
312
        generations[0].generated_text.generated_tokens
313
314
        == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
    )