File size: 1,800 Bytes
0475af5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import pytest
from utils import *

server: ServerProcess


@pytest.fixture(autouse=True)
def create_server():
    global server
    server = ServerPreset.tinyllama2()
    server.gcp_compat = True


def test_gcp_predict_camel_case():
    global server
    server.start()
    res = server.make_request("POST", "/predict", data={
        "instances": [
            {
                "@requestFormat": "chatCompletions",
                "max_tokens": 8,
                "messages": [
                    {"role": "user", "content": "What is the meaning of life?"},
                ],
            }
        ],
    })
    assert res.status_code == 200
    assert "predictions" in res.body
    assert len(res.body["predictions"]) == 1
    prediction = res.body["predictions"][0]
    assert "choices" in prediction
    assert len(prediction["choices"]) == 1
    assert prediction["choices"][0]["message"]["role"] == "assistant"
    assert len(prediction["choices"][0]["message"]["content"]) > 0


def test_gcp_predict_multiple_instances():
    global server
    server.n_slots = 2
    server.start()
    res = server.make_request("POST", "/predict", data={
        "instances": [
            {
                "@requestFormat": "chatCompletions",
                "max_tokens": 8,
                "messages": [{"role": "user", "content": "Say hello"}],
            },
            {
                "@requestFormat": "chatCompletions",
                "max_tokens": 8,
                "messages": [{"role": "user", "content": "Say world"}],
            },
        ],
    })
    assert res.status_code == 200
    assert len(res.body["predictions"]) == 2
    for prediction in res.body["predictions"]:
        assert "choices" in prediction
        assert len(prediction["choices"][0]["message"]["content"]) > 0