text-generation-inference

1
import pytest
2
import torch
3

4
from copy import copy
5
from transformers import AutoTokenizer
6

7
from text_generation_server.pb import generate_pb2
8
from text_generation_server.models.causal_lm import CausalLMBatch
9
from text_generation_server.utils import weight_hub_files, download_weights
10
from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOMSharded
11

12

13
@pytest.fixture(scope="session")
14
def default_bloom():
15
    model_id = "bigscience/bloom-560m"
16
    revision = "main"
17
    filenames = weight_hub_files(model_id, revision, ".safetensors")
18
    download_weights(filenames, model_id, revision)
19
    return BLOOMSharded(model_id)
20

21

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

26

27
@pytest.fixture
28
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
29
    return generate_pb2.Request(
30
        id=0,
31
        inputs="Test",
32
        prefill_logprobs=True,
33
        truncate=100,
34
        parameters=default_pb_parameters,
35
        stopping_parameters=default_pb_stop_parameters,
36
    )
37

38

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

43

44
@pytest.fixture
45
def default_bloom_batch(default_pb_batch, bloom_560m_tokenizer):
46
    return BloomCausalLMBatch.from_pb(
47
        default_pb_batch, bloom_560m_tokenizer, torch.float32, torch.device("cpu")
48
    )
49

50

51
@pytest.fixture
52
def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer):
53
    req_0 = copy(default_pb_request)
54
    req_0.id = 1
55
    req_1 = default_pb_request
56
    req_1.id = 2
57
    req_1.stopping_parameters.max_new_tokens = 5
58

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

64

65
def test_batch_from_pb(default_pb_batch, default_bloom_batch):
66
    batch = default_bloom_batch
67

68
    assert batch.batch_id == default_pb_batch.id
69
    assert batch.requests == default_pb_batch.requests
70

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

75
    assert batch.attention_mask[0][0] == 1
76
    assert torch.all(batch.attention_mask[0][1:] == 0)
77

78
    assert batch.past_key_values is None
79

80
    assert all(
81
        [
82
            torch.equal(input_ids, all_input_ids[:, 0])
83
            for input_ids, all_input_ids in zip(batch.input_ids, batch.all_input_ids)
84
        ]
85
    )
86

87
    assert batch.input_lengths == [1]
88

89
    assert len(batch) == default_pb_batch.size
90
    assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == len(batch)
91

92
    assert batch.max_input_length == batch.input_lengths[0]
93

94

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

99

100
def test_causal_lm_batch_type(default_bloom):
101
    assert default_bloom.batch_type == BloomCausalLMBatch
102

103

104
def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
105
    sequence_length = len(default_bloom_batch.all_input_ids[0])
106
    generations, next_batch, _ = default_bloom.generate_token(default_bloom_batch)
107

108
    assert len(generations) == len(default_bloom_batch)
109
    assert isinstance(next_batch, CausalLMBatch)
110
    assert not next_batch.keys_head_dim_last
111

112
    assert len(next_batch.all_input_ids) == len(next_batch)
113
    assert len(next_batch.all_input_ids[0]) == sequence_length + 1
114
    assert len(next_batch.attention_mask[0]) == 11
115
    assert torch.all(next_batch.all_input_ids[0][-2:] == 10264)
116
    assert torch.all(next_batch.all_input_ids[0][:-2] == 3)
117

118
    assert torch.all(next_batch.attention_mask[0][:2] == 1)
119
    assert torch.all(next_batch.attention_mask[0][2:] == 0)
120

121
    assert next_batch.input_ids.shape == (len(next_batch), 1)
122
    assert next_batch.input_ids[0, 0] == 10264
123

124
    assert next_batch.input_lengths == [2]
125
    assert next_batch.max_input_length == next_batch.input_lengths[0]
126

127
    assert next_batch.past_key_values is not None
128
    assert all(
129
        [p[0].shape == (16, 64, sequence_length) for p in next_batch.past_key_values]
130
    )
131
    assert all(
132
        [p[1].shape == (16, sequence_length, 64) for p in next_batch.past_key_values]
133
    )
134
    assert all([generation.generated_text is None for generation in generations])
135
    assert all([len(generation.prefill_tokens) == 1 for generation in generations])
136
    assert all(
137
        [
138
            token_id.item() == 10264
139
            for generation in generations
140
            for token_id in generation.tokens.token_ids
141
        ]
142
    )
143
    assert all(
144
        [
145
            token_text == "Test"
146
            for generation in generations
147
            for token_text in generation.tokens.texts
148
        ]
149
    )
150
    assert generations[0].request_id == 0
151

152

153
def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch):
154
    next_batch = default_bloom_batch
155
    for _ in range(default_bloom_batch.stopping_criterias[0].max_new_tokens - 1):
156
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
157
        assert len(generations) == len(default_bloom_batch)
158

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

162
    assert len(generations) == 1
163
    assert (
164
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
165
    )
166
    assert generations[0].request_id == default_bloom_batch.requests[0].id
167
    assert (
168
        generations[0].generated_text.generated_tokens
169
        == default_bloom_batch.stopping_criterias[0].max_new_tokens
170
    )
171

172

173
def test_causal_lm_generate_token_completion_multi(
174
    default_bloom, default_multi_requests_bloom_batch
175
):
176
    next_batch = default_multi_requests_bloom_batch
177

178
    for i in range(
179
        default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens - 1
180
    ):
181
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
182
        assert len(generations) == len(default_multi_requests_bloom_batch)
183

184
    generations, next_batch, _ = default_bloom.generate_token(next_batch)
185
    assert next_batch is not None
186

187
    assert len(generations) == 2
188
    assert generations[1].generated_text.text == "TestTestTestTestTest"
189
    assert (
190
        generations[1].request_id == default_multi_requests_bloom_batch.requests[1].id
191
    )
192
    assert (
193
        generations[1].generated_text.generated_tokens
194
        == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
195
    )
196
    # Copy stopping_criterias before filtering
197
    stopping_criterias = default_multi_requests_bloom_batch.stopping_criterias.copy()
198

199
    next_batch = next_batch.filter([next_batch.requests[0].id])
200

201
    for _ in range(
202
        stopping_criterias[0].max_new_tokens - stopping_criterias[1].max_new_tokens - 1
203
    ):
204
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
205
        assert len(generations) == len(next_batch)
206

207
    generations, next_batch, _ = default_bloom.generate_token(next_batch)
208
    assert next_batch is None
209

210
    assert len(generations) == 1
211
    assert (
212
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
213
    )
214
    assert (
215
        generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
216
    )
217
    assert (
218
        generations[0].generated_text.generated_tokens
219
        == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
220
    )
221

222

223
def test_batch_concatenate(
224
    default_bloom, default_bloom_batch, default_multi_requests_bloom_batch
225
):
226
    next_batch_0 = default_bloom_batch
227
    _, next_batch_0, _ = default_bloom.generate_token(next_batch_0)
228
    _, next_batch_0, _ = default_bloom.generate_token(next_batch_0)
229

230
    next_batch_1 = default_multi_requests_bloom_batch
231
    _, next_batch_1, _ = default_bloom.generate_token(next_batch_1)
232

233
    # Clone past_key_values before concatenating to compare after,
234
    # because they are removed from the concatenated batches
235
    next_batch_0_past_key_values = [
236
        (k.clone(), v.clone()) for (k, v) in next_batch_0.past_key_values
237
    ]
238
    next_batch_1_past_key_values = [
239
        (k.clone(), v.clone()) for (k, v) in next_batch_1.past_key_values
240
    ]
241

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

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

248
    assert torch.all(
249
        next_batch.attention_mask[0, : -next_batch.padding_right_offset] == 1
250
    )
251
    assert torch.all(
252
        next_batch.attention_mask[1:, 1 : -next_batch.padding_right_offset] == 1
253
    )
254
    assert torch.all(next_batch.attention_mask[1:, 3:] == 0)
255

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

259
    assert next_batch.input_lengths == [3, 2, 2]
260
    assert next_batch.max_input_length == 3
261

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

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

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

271
    assert next_batch.past_key_values is not None
272
    assert all([p[0].shape == (3, 16, 64, 2) for p in next_batch.past_key_values])
273
    assert all([p[1].shape == (3, 16, 2, 64) for p in next_batch.past_key_values])
274

275
    for i, past in enumerate(next_batch.past_key_values):
276
        assert torch.equal(next_batch_0_past_key_values[i][0][:, :, -2:], past[0][0])
277
        assert torch.equal(
278
            next_batch_1_past_key_values[i][0][:, :, -1:],
279
            past[0][1:, :, :, -1].reshape(-1, 64, 1),
280
        )
281

282
        assert torch.equal(next_batch_0_past_key_values[i][1][:, -2:, :], past[1][0])
283
        assert torch.equal(
284
            next_batch_1_past_key_values[i][1][:, -1:, :],
285
            past[1][1:, :, -1, :].reshape(-1, 1, 64),
286
        )
287

288
    for _ in range(
289
        default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens - 2
290
    ):
291
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
292
        assert len(generations) == len(next_batch)
293

294
    generations, next_batch, _ = default_bloom.generate_token(next_batch)
295
    assert next_batch is not None
296

297
    assert len(generations) == 3
298
    assert generations[2].generated_text.text == "TestTestTestTestTest"
299
    assert (
300
        generations[2].request_id == default_multi_requests_bloom_batch.requests[1].id
301
    )
302
    assert (
303
        generations[2].generated_text.generated_tokens
304
        == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
305
    )
306

307
    next_batch = next_batch.filter(
308
        [next_batch.requests[0].id, next_batch.requests[1].id]
309
    )
310

311
    for _ in range(
312
        default_bloom_batch.stopping_criterias[0].max_new_tokens
313
        - default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
314
        - 2
315
    ):
316
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
317
        assert len(generations) == len(next_batch)
318

319
    generations, next_batch, _ = default_bloom.generate_token(next_batch)
320
    assert next_batch is not None
321

322
    assert len(generations) == 2
323
    assert (
324
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
325
    )
326
    assert generations[0].request_id == default_bloom_batch.requests[0].id
327
    assert (
328
        generations[0].generated_text.generated_tokens
329
        == default_bloom_batch.stopping_criterias[0].max_new_tokens
330
    )
331

332
    next_batch = next_batch.filter([next_batch.requests[1].id])
333

334
    for _ in range(
335
        default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
336
        - default_bloom_batch.stopping_criterias[0].max_new_tokens
337
        - default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
338
        - 4
339
    ):
340
        generations, next_batch, _ = default_bloom.generate_token(next_batch)
341
        assert len(generations) == len(next_batch)
342

343
    generations, next_batch, _ = default_bloom.generate_token(next_batch)
344
    assert next_batch is None
345

346
    assert len(generations) == 1
347
    assert (
348
        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
349
    )
350
    assert (
351
        generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
352
    )
353
    assert (
354
        generations[0].generated_text.generated_tokens
355
        == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
356
    )
357