test_skip_tokenizer_init.py 1.91 KB
Newer Older
1
2
3
4
5
6
import json
import unittest

import requests

from sglang.srt.utils import kill_child_process
7
8
from sglang.test.test_utils import (
    DEFAULT_MODEL_NAME_FOR_TEST,
Yineng Zhang's avatar
Yineng Zhang committed
9
    DEFAULT_URL_FOR_UNIT_TEST,
10
11
    popen_launch_server,
)
12
13


Lianmin Zheng's avatar
Lianmin Zheng committed
14
class TestSkipTokenizerInit(unittest.TestCase):
15
16
17
    @classmethod
    def setUpClass(cls):
        cls.model = DEFAULT_MODEL_NAME_FOR_TEST
Yineng Zhang's avatar
Yineng Zhang committed
18
        cls.base_url = DEFAULT_URL_FOR_UNIT_TEST
19
20
21
22
23
24
25
26
        cls.process = popen_launch_server(
            cls.model, cls.base_url, timeout=300, other_args=["--skip-tokenizer-init"]
        )

    @classmethod
    def tearDownClass(cls):
        kill_child_process(cls.process.pid)

Lianmin Zheng's avatar
Lianmin Zheng committed
27
    def run_decode(self, return_logprob=False, top_logprobs_num=0, n=1):
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
62
        response = requests.post(
            self.base_url + "/generate",
            json={
                "input_ids": [
                    119689,
                    50650,
                    18291,
                    30061,
                    5316,
                    26951,
                    119690,
                ],  # The capital of France is
                "sampling_params": {
                    "temperature": 0 if n == 1 else 0.5,
                    "max_new_tokens": 32,
                    "n": n,
                    "stop_token_ids": [119690],
                },
                "stream": False,
                "return_logprob": return_logprob,
                "top_logprobs_num": top_logprobs_num,
                "logprob_start_len": 0,
            },
        )
        print(json.dumps(response.json()))
        print("=" * 100)

    def test_simple_decode(self):
        self.run_decode()

    def test_parallel_sample(self):
        self.run_decode(n=3)

    def test_logprob(self):
        for top_logprobs_num in [0, 3]:
Lianmin Zheng's avatar
Lianmin Zheng committed
63
64
65
66
            self.run_decode(
                return_logprob=True,
                top_logprobs_num=top_logprobs_num,
            )
67
68
69


if __name__ == "__main__":
Lianmin Zheng's avatar
Lianmin Zheng committed
70
    unittest.main()