api_client.py 2.57 KB
Newer Older
1
"""Example Python client for `vllm.entrypoints.api_server`
2
3
NOTE: The API server is used only for demonstration and simple performance
benchmarks. It is not intended for production use.
4
For production use, we recommend `vllm serve` and the OpenAI client API.
5
"""
6

Zhuohan Li's avatar
Zhuohan Li committed
7
8
import argparse
import json
9
from typing import Iterable, List
Zhuohan Li's avatar
Zhuohan Li committed
10

11
12
13
14
import requests


def clear_line(n: int = 1) -> None:
Zhuohan Li's avatar
Zhuohan Li committed
15
16
    LINE_UP = '\033[1A'
    LINE_CLEAR = '\x1b[2K'
17
    for _ in range(n):
Zhuohan Li's avatar
Zhuohan Li committed
18
19
20
        print(LINE_UP, end=LINE_CLEAR, flush=True)


21
22
23
def post_http_request(prompt: str,
                      api_url: str,
                      n: int = 1,
24
                      stream: bool = False) -> requests.Response:
Zhuohan Li's avatar
Zhuohan Li committed
25
26
27
28
29
30
31
    headers = {"User-Agent": "Test Client"}
    pload = {
        "prompt": prompt,
        "n": n,
        "use_beam_search": True,
        "temperature": 0.0,
        "max_tokens": 16,
32
        "stream": stream,
Zhuohan Li's avatar
Zhuohan Li committed
33
34
    }
    response = requests.post(api_url, headers=headers, json=pload, stream=True)
35
36
    return response

Zhuohan Li's avatar
Zhuohan Li committed
37

38
def get_streaming_response(response: requests.Response) -> Iterable[List[str]]:
39
40
    for chunk in response.iter_lines(chunk_size=8192,
                                     decode_unicode=False,
41
                                     delimiter=b"\0"):
Zhuohan Li's avatar
Zhuohan Li committed
42
43
44
45
46
47
        if chunk:
            data = json.loads(chunk.decode("utf-8"))
            output = data["text"]
            yield output


48
49
50
51
52
53
def get_response(response: requests.Response) -> List[str]:
    data = json.loads(response.content)
    output = data["text"]
    return output


Zhuohan Li's avatar
Zhuohan Li committed
54
55
56
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default="localhost")
57
    parser.add_argument("--port", type=int, default=8000)
Zhuohan Li's avatar
Zhuohan Li committed
58
59
    parser.add_argument("--n", type=int, default=4)
    parser.add_argument("--prompt", type=str, default="San Francisco is a")
60
    parser.add_argument("--stream", action="store_true")
Zhuohan Li's avatar
Zhuohan Li committed
61
62
63
64
    args = parser.parse_args()
    prompt = args.prompt
    api_url = f"http://{args.host}:{args.port}/generate"
    n = args.n
65
    stream = args.stream
Zhuohan Li's avatar
Zhuohan Li committed
66

67
    print(f"Prompt: {prompt!r}\n", flush=True)
68
69
70
    response = post_http_request(prompt, api_url, n, stream)

    if stream:
Zhuohan Li's avatar
Zhuohan Li committed
71
        num_printed_lines = 0
72
73
74
75
76
        for h in get_streaming_response(response):
            clear_line(num_printed_lines)
            num_printed_lines = 0
            for i, line in enumerate(h):
                num_printed_lines += 1
77
                print(f"Beam candidate {i}: {line!r}", flush=True)
78
79
80
    else:
        output = get_response(response)
        for i, line in enumerate(output):
81
            print(f"Beam candidate {i}: {line!r}", flush=True)