test_causal_lm.py 12.2 KB
Newer Older
1
2
3
4
import pytest
import torch

from copy import copy
5
from transformers import AutoTokenizer
6

7
8
from text_generation_server.pb import generate_pb2
from text_generation_server.models.causal_lm import CausalLM, CausalLMBatch
9
10


11
12
13
14
15
16
17
18
19
20
21
22
@pytest.fixture(scope="session")
def default_causal_lm():
    return CausalLM("gpt2")


@pytest.fixture(scope="session")
def gpt2_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained("gpt2", padding_side="left")
    tokenizer.pad_token_id = 50256
    return tokenizer


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


@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_causal_lm_batch(default_pb_batch, gpt2_tokenizer):
42
43
44
    return CausalLMBatch.from_pb(
        default_pb_batch, gpt2_tokenizer, torch.float32, torch.device("cpu")
    )
45
46
47
48
49


@pytest.fixture
def default_multi_requests_causal_lm_batch(default_pb_request, gpt2_tokenizer):
    req_0 = copy(default_pb_request)
50
    req_0.id = 1
51
    req_1 = default_pb_request
52
    req_1.id = 2
53
    req_1.stopping_parameters.max_new_tokens = 5
54

55
    batch_pb = generate_pb2.Batch(id=1, requests=[req_0, req_1], size=2)
56
57
58
    return CausalLMBatch.from_pb(
        batch_pb, gpt2_tokenizer, torch.float32, torch.device("cpu")
    )
59
60
61
62
63
64
65
66
67
68
69
70


def test_batch_from_pb(default_pb_batch, default_causal_lm_batch):
    batch = default_causal_lm_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] == 14402
    assert torch.all(batch.input_ids[0][:-1] == 50256)

71
72
    assert batch.attention_mask[0, 0] == 1
    assert torch.all(batch.attention_mask[0, 1:] == 0)
73
74
75

    assert batch.past_key_values is None

76
77
78
79
80
81
    assert all(
        [
            torch.equal(input_ids, all_input_ids[:, 0])
            for input_ids, all_input_ids in zip(batch.input_ids, batch.all_input_ids)
        ]
    )
82
83
84

    assert batch.input_lengths == [1]

85
86
    assert len(batch) == default_pb_batch.size
    assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == len(batch)
87

88
    assert batch.max_input_length == batch.input_lengths[0]
89
90
91
92
93
94
95
96
97
98
99
100


def test_batch_concatenate_no_prefill(default_causal_lm_batch):
    with pytest.raises(ValueError):
        CausalLMBatch.concatenate([default_causal_lm_batch, default_causal_lm_batch])


def test_causal_lm_batch_type(default_causal_lm):
    assert default_causal_lm.batch_type == CausalLMBatch


def test_causal_lm_generate_token(default_causal_lm, default_causal_lm_batch):
101
    sequence_length = len(default_causal_lm_batch.all_input_ids[0])
102
    generations, next_batch = default_causal_lm.generate_token(default_causal_lm_batch)
103

104
    assert len(generations) == len(next_batch)
105
106
    assert isinstance(next_batch, CausalLMBatch)

107
    assert len(next_batch.all_input_ids) == len(next_batch)
108
109
    assert len(next_batch.all_input_ids[0]) == sequence_length + 1
    assert len(next_batch.attention_mask[0]) == 11
110
    assert next_batch.all_input_ids[0][-1] == 13
111
112
113
    assert next_batch.all_input_ids[0][-2] == 14402
    assert torch.all(next_batch.all_input_ids[0][:-2] == 50256)

114
115
    assert torch.all(next_batch.attention_mask[0][0:2] == 1)
    assert torch.all(next_batch.attention_mask[0][2:] == 0)
116

117
    assert next_batch.input_ids.shape == (len(next_batch), 1)
118
    assert next_batch.input_ids[0, 0] == 13
119
120

    assert next_batch.input_lengths == [2]
121
    assert next_batch.max_input_length == next_batch.input_lengths[0]
122
123

    assert next_batch.past_key_values is not None
124
125
126
127
128
129
    assert all(
        [p[0].shape == (1, 12, sequence_length, 64) for p in next_batch.past_key_values]
    )
    assert all(
        [p[1].shape == (1, 12, sequence_length, 64) for p in next_batch.past_key_values]
    )
130
131
132
133
134
    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() == 13 for generation in generations])
    assert all([generation.token_text == "." for generation in generations])
    assert generations[0].request_id == 0
135
136
137
138
139
140
141


def test_causal_lm_generate_token_completion(
    default_causal_lm, default_causal_lm_batch
):
    next_batch = default_causal_lm_batch
    for _ in range(default_causal_lm_batch.stopping_criterias[0].max_new_tokens - 1):
142
143
        generations, next_batch = default_causal_lm.generate_token(next_batch)
        assert len(generations) == len(next_batch)
144

145
    generations, next_batch = default_causal_lm.generate_token(next_batch)
146
147
    assert next_batch is None

148
    assert len(generations) == 1
149
    assert generations[0].generated_text.text == ".java:784) at net.minecraft."
150
    assert generations[0].request_id == default_causal_lm_batch.requests[0].id
151
    assert (
152
        generations[0].generated_text.generated_tokens
153
154
155
156
157
158
159
160
161
162
163
164
        == default_causal_lm_batch.stopping_criterias[0].max_new_tokens
    )


def test_causal_lm_generate_token_completion_multi(
    default_causal_lm, default_multi_requests_causal_lm_batch
):
    next_batch = default_multi_requests_causal_lm_batch

    for i in range(
        default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens - 1
    ):
165
166
        generations, next_batch = default_causal_lm.generate_token(next_batch)
        assert len(generations) == len(next_batch)
167

168
    generations, next_batch = default_causal_lm.generate_token(next_batch)
169
170
    assert next_batch is not None

171
    assert len(generations) == 2
172
    assert generations[1].generated_text.text == ".java:784)"
173
    assert (
174
175
        generations[1].request_id
        == default_multi_requests_causal_lm_batch.requests[1].id
176
177
    )
    assert (
178
        generations[1].generated_text.generated_tokens
179
180
        == default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
    )
181
    # Copy stopping_criterias before filtering
182
183
184
    stopping_criterias = (
        default_multi_requests_causal_lm_batch.stopping_criterias.copy()
    )
185

186
    next_batch = next_batch.filter([next_batch.requests[0].id])
187

188
    for _ in range(
189
        stopping_criterias[0].max_new_tokens - stopping_criterias[1].max_new_tokens - 1
190
    ):
191
192
        generations, next_batch = default_causal_lm.generate_token(next_batch)
        assert len(generations) == len(next_batch)
193

194
    generations, next_batch = default_causal_lm.generate_token(next_batch)
195
196
    assert next_batch is None

197
    assert len(generations) == 1
198
    assert generations[0].generated_text.text == ".java:784) at net.minecraft."
199
    assert (
200
201
        generations[0].request_id
        == default_multi_requests_causal_lm_batch.requests[0].id
202
203
    )
    assert (
204
        generations[0].generated_text.generated_tokens
205
206
207
208
209
210
211
212
213
214
215
216
217
218
        == default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
    )


def test_batch_concatenate(
    default_causal_lm, default_causal_lm_batch, default_multi_requests_causal_lm_batch
):
    next_batch_0 = default_causal_lm_batch
    _, next_batch_0 = default_causal_lm.generate_token(next_batch_0)
    _, next_batch_0 = default_causal_lm.generate_token(next_batch_0)

    next_batch_1 = default_multi_requests_causal_lm_batch
    _, next_batch_1 = default_causal_lm.generate_token(next_batch_1)

219
220
221
222
223
224
225
226
227
    # Clone past_key_values before concatenating to compare after,
    # because they are removed from the concatenated batches
    next_batch_0_past_key_values = [
        (k.clone(), v.clone()) for (k, v) in next_batch_0.past_key_values
    ]
    next_batch_1_past_key_values = [
        (k.clone(), v.clone()) for (k, v) in next_batch_1.past_key_values
    ]

228
229
230
231
232
233
    next_batch = CausalLMBatch.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])

234
235
236
237
238
239
240
    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)
241
242

    assert next_batch.batch_id == 0
243
244
    assert next_batch.input_ids[0, 0] == 12355
    assert torch.all(next_batch.input_ids[1:] == 13)
245
246

    assert next_batch.input_lengths == [3, 2, 2]
247
    assert next_batch.max_input_length == 3
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

    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, 12, 2, 64) for p in next_batch.past_key_values])
    assert all([p[1].shape == (3, 12, 2, 64) for p in next_batch.past_key_values])

    for i, past in enumerate(next_batch.past_key_values):
263
        assert torch.equal(next_batch_0_past_key_values[i][0][0, :, -2:], past[0][0])
264
        assert torch.equal(
265
            next_batch_1_past_key_values[i][0][:, :, -1:], past[0][1:, :, -1:, :]
266
267
        )

268
        assert torch.equal(next_batch_0_past_key_values[i][1][0, :, -2:], past[1][0])
269
        assert torch.equal(
270
            next_batch_1_past_key_values[i][1][:, :, -1:], past[1][1:, :, -1:, :]
271
272
273
274
275
        )

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

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

282
    assert len(generations) == 3
283
    assert generations[2].generated_text.text == ".java:784)"
284
    assert (
285
286
        generations[2].request_id
        == default_multi_requests_causal_lm_batch.requests[1].id
287
288
    )
    assert (
289
        generations[2].generated_text.generated_tokens
290
291
292
        == default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
    )

293
294
295
    next_batch = next_batch.filter(
        [next_batch.requests[0].id, next_batch.requests[1].id]
    )
296

297
298
299
300
301
    for _ in range(
        default_causal_lm_batch.stopping_criterias[0].max_new_tokens
        - default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
        - 2
    ):
302
303
        generations, next_batch = default_causal_lm.generate_token(next_batch)
        assert len(generations) == len(next_batch)
304

305
    generations, next_batch = default_causal_lm.generate_token(next_batch)
306
307
    assert next_batch is not None

308
    assert len(generations) == 2
309
    assert generations[0].generated_text.text == ".java:784) at net.minecraft."
310
    assert generations[0].request_id == default_causal_lm_batch.requests[0].id
311
    assert (
312
        generations[0].generated_text.generated_tokens
313
314
315
        == default_causal_lm_batch.stopping_criterias[0].max_new_tokens
    )

316
    next_batch = next_batch.filter([next_batch.requests[1].id])
317

318
319
320
321
322
323
    for _ in range(
        default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
        - default_causal_lm_batch.stopping_criterias[0].max_new_tokens
        - default_multi_requests_causal_lm_batch.stopping_criterias[1].max_new_tokens
        - 4
    ):
324
325
        generations, next_batch = default_causal_lm.generate_token(next_batch)
        assert len(generations) == len(next_batch)
326

327
    generations, next_batch = default_causal_lm.generate_token(next_batch)
328
329
    assert next_batch is None

330
    assert len(generations) == 1
331
    assert generations[0].generated_text.text == ".java:784) at net.minecraft."
332
    assert (
333
334
        generations[0].request_id
        == default_multi_requests_causal_lm_batch.requests[0].id
335
336
    )
    assert (
337
        generations[0].generated_text.generated_tokens
338
339
        == default_multi_requests_causal_lm_batch.stopping_criterias[0].max_new_tokens
    )