text-generation-inference

1
import pytest
2

3
from text_generation_server.pb import generate_pb2
4
from text_generation_server.models.causal_lm import CausalLMBatch
5
from text_generation_server.models.santacoder import SantaCoder
6

7

8
@pytest.fixture(scope="session")
9
def default_santacoder():
10
    return SantaCoder("bigcode/santacoder")
11

12

13
@pytest.fixture
14
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
15
    return generate_pb2.Request(
16
        id=0,
17
        inputs="def",
18
        prefill_logprobs=True,
19
        truncate=100,
20
        parameters=default_pb_parameters,
21
        stopping_parameters=default_pb_stop_parameters,
22
    )
23

24

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

29

30
@pytest.fixture
31
def default_fim_pb_request(default_pb_parameters, default_pb_stop_parameters):
32
    return generate_pb2.Request(
33
        id=0,
34
        inputs="<fim-prefix>def<fim-suffix>world<fim-middle>",
35
        prefill_logprobs=True,
36
        truncate=100,
37
        parameters=default_pb_parameters,
38
        stopping_parameters=default_pb_stop_parameters,
39
    )
40

41

42
@pytest.fixture
43
def default_fim_pb_batch(default_fim_pb_request):
44
    return generate_pb2.Batch(id=0, requests=[default_fim_pb_request], size=1)
45

46

47
@pytest.mark.skip
48
def test_santacoder_generate_token_completion(default_santacoder, default_pb_batch):
49
    batch = CausalLMBatch.from_pb(
50
        default_pb_batch,
51
        default_santacoder.tokenizer,
52
        default_santacoder.dtype,
53
        default_santacoder.device,
54
    )
55
    next_batch = batch
56

57
    for _ in range(batch.stopping_criterias[0].max_new_tokens - 1):
58
        generations, next_batch, _ = default_santacoder.generate_token(next_batch)
59
        assert len(generations) == len(next_batch)
60

61
    generations, next_batch, _ = default_santacoder.generate_token(next_batch)
62
    assert next_batch is None
63

64
    assert len(generations) == 1
65
    assert generations[0].generated_text.text == " test_get_all_users_with_"
66
    assert generations[0].request_id == batch.requests[0].id
67
    assert (
68
        generations[0].generated_text.generated_tokens
69
        == batch.stopping_criterias[0].max_new_tokens
70
    )
71

72

73
@pytest.mark.skip
74
def test_fim_santacoder_generate_token_completion(
75
    default_santacoder, default_fim_pb_batch
76
):
77
    batch = CausalLMBatch.from_pb(
78
        default_fim_pb_batch,
79
        default_santacoder.tokenizer,
80
        default_santacoder.dtype,
81
        default_santacoder.device,
82
    )
83
    next_batch = batch
84

85
    for _ in range(batch.stopping_criterias[0].max_new_tokens - 1):
86
        generations, next_batch, _ = default_santacoder.generate_token(next_batch)
87
        assert len(generations) == len(next_batch)
88

89
    generations, next_batch, _ = default_santacoder.generate_token(next_batch)
90
    assert next_batch is None
91

92
    assert len(generations) == 1
93
    assert (
94
        generations[0].generated_text.text
95
        == """ineProperty(exports, "__esModule", { value"""
96
    )
97
    assert generations[0].request_id == batch.requests[0].id
98
    assert (
99
        generations[0].generated_text.generated_tokens
100
        == batch.stopping_criterias[0].max_new_tokens
101
    )
102