main
py 75 lines 2.22 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Example for starting a Gradio Webserver
4 Start vLLM API server:
5 python examples/applications/api_server/server.py \
6 --model meta-llama/Llama-2-7b-chat-hf
7
8 Start Webserver:
9 python examples/applications/chatbot/gradio_webserver.py
10
11 Note that `pip install --upgrade gradio` is needed to run this example.
12 More details: https://github.com/gradio-app/gradio
13
14 If your antivirus software blocks the download of frpc for gradio,
15 you can install it manually by following these steps:
16
17 1. Download this file: https://cdn-media.huggingface.co/frpc-gradio-0.3/frpc_linux_amd64
18 2. Rename the downloaded file to: frpc_linux_amd64_v0.3
19 3. Move the file to this location: /home/user/.cache/huggingface/gradio/frpc
20 """
21
22 import argparse
23 import json
24
25 import gradio as gr
26 import requests
27
28
29 def http_bot(prompt):
30 headers = {"User-Agent": "vLLM Client"}
31 pload = {
32 "prompt": prompt,
33 "stream": True,
34 "max_tokens": 128,
35 }
36 response = requests.post(args.model_url, headers=headers, json=pload, stream=True)
37
38 for chunk in response.iter_lines(
39 chunk_size=8192, decode_unicode=False, delimiter=b"\n"
40 ):
41 if chunk:
42 data = json.loads(chunk.decode("utf-8"))
43 output = data["text"][0]
44 yield output
45
46
47 def build_demo():
48 with gr.Blocks() as demo:
49 gr.Markdown("# vLLM text completion demo\n")
50 inputbox = gr.Textbox(label="Input", placeholder="Enter text and press ENTER")
51 outputbox = gr.Textbox(
52 label="Output", placeholder="Generated result from the model"
53 )
54 inputbox.submit(http_bot, [inputbox], [outputbox])
55 return demo
56
57
58 def parse_args():
59 parser = argparse.ArgumentParser()
60 parser.add_argument("--host", type=str, default=None)
61 parser.add_argument("--port", type=int, default=8001)
62 parser.add_argument(
63 "--model-url", type=str, default="http://localhost:8000/generate"
64 )
65 return parser.parse_args()
66
67
68 def main(args):
69 demo = build_demo()
70 demo.queue().launch(server_name=args.host, server_port=args.port, share=True)
71
72
73 if __name__ == "__main__":
74 args = parse_args()
75 main(args)