364 lines
13 KiB
Python
364 lines
13 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from copy import copy
|
|
from transformers import AutoTokenizer
|
|
|
|
from text_generation_server.pb import generate_pb2
|
|
from text_generation_server.models.causal_lm import CausalLMBatch
|
|
from text_generation_server.utils import weight_hub_files, download_weights
|
|
from text_generation_server.models.bloom import BloomCausalLMBatch, BLOOMSharded
|
|
from text_generation_server.models.custom_modeling.bloom_modeling import (
|
|
BloomForCausalLM,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def default_bloom():
|
|
model_id = "bigscience/bloom-560m"
|
|
revision = "main"
|
|
filenames = weight_hub_files(model_id, revision, ".safetensors")
|
|
download_weights(filenames, model_id, revision)
|
|
return BLOOMSharded(
|
|
model_id,
|
|
model_class=BloomForCausalLM,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def bloom_560m_tokenizer():
|
|
return AutoTokenizer.from_pretrained("bigscience/bloom-560m", padding_side="left")
|
|
|
|
|
|
@pytest.fixture
|
|
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
|
|
return generate_pb2.Request(
|
|
id=0,
|
|
inputs="Test",
|
|
input_chunks=generate_pb2.Input(chunks=[generate_pb2.InputChunk(text="Test")]),
|
|
prefill_logprobs=True,
|
|
truncate=100,
|
|
parameters=default_pb_parameters,
|
|
stopping_parameters=default_pb_stop_parameters,
|
|
)
|
|
|
|
|
|
@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.float32, torch.device("cpu")
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer):
|
|
req_0 = copy(default_pb_request)
|
|
req_0.id = 1
|
|
req_1 = default_pb_request
|
|
req_1.id = 2
|
|
req_1.stopping_parameters.max_new_tokens = 5
|
|
|
|
batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
|
|
return BloomCausalLMBatch.from_pb(
|
|
batch_pb, bloom_560m_tokenizer, torch.float32, 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)
|
|
|
|
assert batch.attention_mask[0][0] == 1
|
|
assert torch.all(batch.attention_mask[0][1:] == 0)
|
|
|
|
assert batch.past_key_values is None
|
|
|
|
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)
|
|
]
|
|
)
|
|
|
|
assert batch.input_lengths == [1]
|
|
|
|
assert len(batch) == default_pb_batch.size
|
|
assert len(batch.next_token_choosers) == len(batch.stopping_criterias) == len(batch)
|
|
|
|
assert batch.max_input_length == batch.input_lengths[0]
|
|
|
|
|
|
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):
|
|
sequence_length = len(default_bloom_batch.all_input_ids[0])
|
|
generations, next_batch, _ = default_bloom.generate_token(default_bloom_batch)
|
|
|
|
assert len(generations) == len(default_bloom_batch)
|
|
assert isinstance(next_batch, CausalLMBatch)
|
|
assert not next_batch.keys_head_dim_last
|
|
|
|
assert len(next_batch.all_input_ids) == len(next_batch)
|
|
assert len(next_batch.all_input_ids[0]) == sequence_length + 1
|
|
assert len(next_batch.attention_mask[0]) == 11
|
|
assert torch.all(next_batch.all_input_ids[0][-2:] == 10264)
|
|
assert torch.all(next_batch.all_input_ids[0][:-2] == 3)
|
|
|
|
assert torch.all(next_batch.attention_mask[0][:2] == 1)
|
|
assert torch.all(next_batch.attention_mask[0][2:] == 0)
|
|
|
|
assert next_batch.input_ids.shape == (len(next_batch), 1)
|
|
assert next_batch.input_ids[0, 0] == 10264
|
|
|
|
assert next_batch.input_lengths == [2]
|
|
assert next_batch.max_input_length == next_batch.input_lengths[0]
|
|
|
|
assert next_batch.past_key_values is not None
|
|
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]
|
|
)
|
|
assert all([generation.generated_text is None for generation in generations])
|
|
assert all([len(generation.prefill_tokens) == 1 for generation in generations])
|
|
assert all(
|
|
[
|
|
token_id.item() == 10264
|
|
for generation in generations
|
|
for token_id in generation.tokens.token_ids
|
|
]
|
|
)
|
|
assert all(
|
|
[
|
|
token_text == "Test"
|
|
for generation in generations
|
|
for token_text in generation.tokens.texts
|
|
]
|
|
)
|
|
assert generations[0].request_id == 0
|
|
|
|
|
|
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):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(default_bloom_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is None
|
|
|
|
assert len(generations) == 1
|
|
assert (
|
|
generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
|
|
)
|
|
assert generations[0].request_id == default_bloom_batch.requests[0].id
|
|
assert (
|
|
generations[0].generated_text.generated_tokens
|
|
== 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
|
|
):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(default_multi_requests_bloom_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is not None
|
|
|
|
assert len(generations) == 2
|
|
assert generations[1].generated_text.text == "TestTestTestTestTest"
|
|
assert (
|
|
generations[1].request_id == default_multi_requests_bloom_batch.requests[1].id
|
|
)
|
|
assert (
|
|
generations[1].generated_text.generated_tokens
|
|
== default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
|
|
)
|
|
# Copy stopping_criterias before filtering
|
|
stopping_criterias = default_multi_requests_bloom_batch.stopping_criterias.copy()
|
|
|
|
next_batch = next_batch.filter([next_batch.requests[0].id])
|
|
|
|
for _ in range(
|
|
stopping_criterias[0].max_new_tokens - stopping_criterias[1].max_new_tokens - 1
|
|
):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(next_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is None
|
|
|
|
assert len(generations) == 1
|
|
assert (
|
|
generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
|
|
)
|
|
assert (
|
|
generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
|
|
)
|
|
assert (
|
|
generations[0].generated_text.generated_tokens
|
|
== 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)
|
|
|
|
# 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
|
|
]
|
|
|
|
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])
|
|
|
|
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)
|
|
|
|
assert next_batch.batch_id == 0
|
|
assert torch.all(next_batch.input_ids == 10264)
|
|
|
|
assert next_batch.input_lengths == [3, 2, 2]
|
|
assert next_batch.max_input_length == 3
|
|
|
|
assert next_batch.requests[0] == next_batch_0.requests[0]
|
|
assert next_batch.requests[1:] == list(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
|
|
):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(next_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is not None
|
|
|
|
assert len(generations) == 3
|
|
assert generations[2].generated_text.text == "TestTestTestTestTest"
|
|
assert (
|
|
generations[2].request_id == default_multi_requests_bloom_batch.requests[1].id
|
|
)
|
|
assert (
|
|
generations[2].generated_text.generated_tokens
|
|
== default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
|
|
)
|
|
|
|
next_batch = next_batch.filter(
|
|
[next_batch.requests[0].id, next_batch.requests[1].id]
|
|
)
|
|
|
|
for _ in range(
|
|
default_bloom_batch.stopping_criterias[0].max_new_tokens
|
|
- default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
|
|
- 2
|
|
):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(next_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is not None
|
|
|
|
assert len(generations) == 2
|
|
assert (
|
|
generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
|
|
)
|
|
assert generations[0].request_id == default_bloom_batch.requests[0].id
|
|
assert (
|
|
generations[0].generated_text.generated_tokens
|
|
== default_bloom_batch.stopping_criterias[0].max_new_tokens
|
|
)
|
|
|
|
next_batch = next_batch.filter([next_batch.requests[1].id])
|
|
|
|
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
|
|
):
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert len(generations) == len(next_batch)
|
|
|
|
generations, next_batch, _ = default_bloom.generate_token(next_batch)
|
|
assert next_batch is None
|
|
|
|
assert len(generations) == 1
|
|
assert (
|
|
generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
|
|
)
|
|
assert (
|
|
generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
|
|
)
|
|
assert (
|
|
generations[0].generated_text.generated_tokens
|
|
== default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
|
|
)
|