test_santacoder.py 3.35 KB
Newer Older
1
2
import pytest

3
from text_generation_server.pb import generate_pb2
4
from text_generation_server.models.causal_lm import CausalLMBatch, CausalLM
5
6
7
8


@pytest.fixture(scope="session")
def default_santacoder():
9
    return CausalLM.fallback(model_id="bigcode/santacoder")
10
11
12
13
14
15
16


@pytest.fixture
def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
    return generate_pb2.Request(
        id=0,
        inputs="def",
Daniël de Kok's avatar
Daniël de Kok committed
17
        input_chunks=generate_pb2.Input(chunks=[generate_pb2.InputChunk(text="def")]),
18
        prefill_logprobs=True,
19
        truncate=100,
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
        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_fim_pb_request(default_pb_parameters, default_pb_stop_parameters):
    return generate_pb2.Request(
        id=0,
        inputs="<fim-prefix>def<fim-suffix>world<fim-middle>",
Daniël de Kok's avatar
Daniël de Kok committed
35
36
37
38
39
40
41
        input_chunks=generate_pb2.Input(
            chunks=[
                generate_pb2.InputChunk(
                    text="<fim-prefix>def<fim-suffix>world<fim-middle>"
                )
            ]
        ),
42
        prefill_logprobs=True,
43
        truncate=100,
44
45
46
47
48
49
50
51
52
53
        parameters=default_pb_parameters,
        stopping_parameters=default_pb_stop_parameters,
    )


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


54
@pytest.mark.skip
55
56
def test_santacoder_generate_token_completion(default_santacoder, default_pb_batch):
    batch = CausalLMBatch.from_pb(
57
58
59
60
        default_pb_batch,
        default_santacoder.tokenizer,
        default_santacoder.dtype,
        default_santacoder.device,
61
62
63
64
    )
    next_batch = batch

    for _ in range(batch.stopping_criterias[0].max_new_tokens - 1):
65
        generations, next_batch, _ = default_santacoder.generate_token(next_batch)
66
        assert len(generations) == len(next_batch)
67

68
    generations, next_batch, _ = default_santacoder.generate_token(next_batch)
69
70
    assert next_batch is None

71
    assert len(generations) == 1
72
    assert generations[0].generated_text.text == " test_get_all_users_with_"
73
    assert generations[0].request_id == batch.requests[0].id
74
    assert (
75
        generations[0].generated_text.generated_tokens
76
77
78
79
        == batch.stopping_criterias[0].max_new_tokens
    )


80
@pytest.mark.skip
81
82
83
84
def test_fim_santacoder_generate_token_completion(
    default_santacoder, default_fim_pb_batch
):
    batch = CausalLMBatch.from_pb(
85
86
87
88
        default_fim_pb_batch,
        default_santacoder.tokenizer,
        default_santacoder.dtype,
        default_santacoder.device,
89
90
91
92
    )
    next_batch = batch

    for _ in range(batch.stopping_criterias[0].max_new_tokens - 1):
93
        generations, next_batch, _ = default_santacoder.generate_token(next_batch)
94
        assert len(generations) == len(next_batch)
95

96
    generations, next_batch, _ = default_santacoder.generate_token(next_batch)
97
98
    assert next_batch is None

99
    assert len(generations) == 1
100
    assert (
101
        generations[0].generated_text.text
102
        == """ineProperty(exports, "__esModule", { value"""
103
    )
104
    assert generations[0].request_id == batch.requests[0].id
105
    assert (
106
        generations[0].generated_text.generated_tokens
107
108
        == batch.stopping_criterias[0].max_new_tokens
    )