github
/
ChatGLM-6B
mirror of https://github.com/THUDM/ChatGLM-6B
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Update api.py 

简洁一些,用库里的函数to_thread
pull/1340/head
aleimu 2023-07-18 19:19:26 +08:00 committed by GitHub
parent 044cf323a7
commit 6ba39b421c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
api.py
View File

@ -1,9 +1,8 @@
import torch
import asyncio
import concurrent.futures
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModel from transformers import AutoTokenizer, AutoModel
import uvicorn, json, datetime import uvicorn, json, datetime
import torch
import asyncio
DEVICE = "cuda" DEVICE = "cuda"
DEVICE_ID = "0" DEVICE_ID = "0"
@ -52,10 +51,7 @@ async def create_item(request: Request):
torch_gc() torch_gc()
return answer return answer
loop = asyncio.get_event_loop() return await asyncio.to_thread(_sync_chat, history=history)
executor = concurrent.futures.ThreadPoolExecutor()
answer = await loop.run_in_executor(executor, _sync_chat, history)
return answer
if __name__ == '__main__': if __name__ == '__main__':