client.py 742 Bytes
Newer Older
Rayyyyy's avatar
Rayyyyy committed
1
2
3
4
5
6
import json
import argparse
import requests

parse = argparse.ArgumentParser()
parse.add_argument('--query', default='请写一首诗')
7
parse.add_argument('--use_hf', action='store_true')
Rayyyyy's avatar
Rayyyyy committed
8
9
10
11
12
13
14
15
16
17
args = parse.parse_args()

print(args.query)
headers = {"Content-Type": "application/json"}
data = {
        "query": args.query,
        "history": []
        }

json_str = json.dumps(data)
18
19
20
21
if args.use_hf:
    response = requests.post("http://localhost:8888/hf_inference", headers=headers, data=json_str.encode("utf-8"), verify=False)
else:
    response = requests.post("http://localhost:8888/vllm_inference", headers=headers, data=json_str.encode("utf-8"), verify=False)
Rayyyyy's avatar
Rayyyyy committed
22
23
24

str_response = response.content.decode("utf-8")
print(json.loads(str_response))