Skip to content

Commit

Permalink
add api benchmark script.
Browse files Browse the repository at this point in the history
  • Loading branch information
is committed May 14, 2023
1 parent 186d656 commit 7e335cc
Showing 1 changed file with 41 additions and 0 deletions.
41 changes: 41 additions & 0 deletions contrib/benchmark.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
import requests
import threading
import logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s.%(msecs)03d %(levelname)s %(module)s - %(funcName)s: %(message)s',
datefmt='%Y-%m-%d %H:%M:%S',
)

def req_model(t_index, request_str):
logging.info(f'start {t_index}:{request_str}')
json={'prompt':request_str,'history':[["现在你是一只小狗","好的,你想聊些什么,汪!"],]}
headers={'Content-Type':'application/json'}
for i in range(3):
try:
r = requests.post(
'http://127.0.0.1:8000',
json=json,
headers=headers,
timeout=60
)
logging.info(f'end {t_index}:{r.json()["response"]}')
return
except:
logging.info(f'end {t_index}:timeout, try again.')
logging.error(f'end {t_index}:timeout 3 times.')

if __name__ == "__main__":
threads=[]
for i in range(10):
threads.append(threading.Thread(target=req_model, args=(i,"你是谁")))
for i in range(10):
threads.append(threading.Thread(target=req_model, args=(i,"讲个笑话吧")))
for t in threads:
t.start()
for t in threads:
t.join()
logging.info('Done.')



0 comments on commit 7e335cc

Please sign in to comment.