From 9cc1bd5136a1544b9171b1d276eb2ba362ab5e6c Mon Sep 17 00:00:00 2001
From: Yu Xin <scaner@gmail.com>
Date: Sun, 14 May 2023 23:19:08 +0800
Subject: [PATCH] =?UTF-8?q?api.py=E6=94=AF=E6=8C=81=E5=B9=B6=E5=8F=91?=
 =?UTF-8?q?=E6=9C=8D=E5=8A=A1.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 api.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/api.py b/api.py
index 693c70a..3193d2d 100644
--- a/api.py
+++ b/api.py
@@ -2,6 +2,8 @@ from fastapi import FastAPI, Request
 from transformers import AutoTokenizer, AutoModel
 import uvicorn, json, datetime
 import torch
+import threading
+import asyncio
 
 DEVICE = "cuda"
 DEVICE_ID = "0"
@@ -17,9 +19,12 @@ def torch_gc():
 
 app = FastAPI()
 
+import concurrent
+from functools import partial
+pool = concurrent.futures.ThreadPoolExecutor(10)
 
 @app.post("/")
-async def create_item(request: Request):
+async def _create_item(request: Request):
     global model, tokenizer
     json_post_raw = await request.json()
     json_post = json.dumps(json_post_raw)
@@ -29,12 +34,13 @@ async def create_item(request: Request):
     max_length = json_post_list.get('max_length')
     top_p = json_post_list.get('top_p')
     temperature = json_post_list.get('temperature')
-    response, history = model.chat(tokenizer,
+    loop = asyncio.get_event_loop()
+    response, history = await loop.run_in_executor(pool,partial(model.chat,tokenizer,
                                    prompt,
                                    history=history,
                                    max_length=max_length if max_length else 2048,
                                    top_p=top_p if top_p else 0.7,
-                                   temperature=temperature if temperature else 0.95)
+                                   temperature=temperature if temperature else 0.95))
     now = datetime.datetime.now()
     time = now.strftime("%Y-%m-%d %H:%M:%S")
     answer = {
@@ -48,9 +54,16 @@ async def create_item(request: Request):
     torch_gc()
     return answer
 
+async def create_item(request: Request):
+    loop = asyncio.get_event_loop()
+    with concurrent.futures.ThreadPoolExecutor() as pool:
+        result = await loop.run_in_executor(pool,_create_item, request)
+        print(result)
+    return result
 
 if __name__ == '__main__':
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
     model.eval()
     uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)
+