pull/1340/merge
aleimu 2024-06-28 19:41:47 +08:00 committed by GitHub
commit f0f9113fa5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 23 additions and 18 deletions

41
api.py
View File

@ -2,6 +2,7 @@ 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 torch
import asyncio
DEVICE = "cuda" DEVICE = "cuda"
DEVICE_ID = "0" DEVICE_ID = "0"
@ -29,24 +30,28 @@ async def create_item(request: Request):
max_length = json_post_list.get('max_length') max_length = json_post_list.get('max_length')
top_p = json_post_list.get('top_p') top_p = json_post_list.get('top_p')
temperature = json_post_list.get('temperature') temperature = json_post_list.get('temperature')
response, history = model.chat(tokenizer,
prompt, def _sync_chat(history):
history=history, response, history = model.chat(tokenizer,
max_length=max_length if max_length else 2048, prompt,
top_p=top_p if top_p else 0.7, history=history,
temperature=temperature if temperature else 0.95) max_length=max_length if max_length else 2048,
now = datetime.datetime.now() top_p=top_p if top_p else 0.7,
time = now.strftime("%Y-%m-%d %H:%M:%S") temperature=temperature if temperature else 0.95)
answer = { now = datetime.datetime.now()
"response": response, time = now.strftime("%Y-%m-%d %H:%M:%S")
"history": history, answer = {
"status": 200, "response": response,
"time": time "history": history,
} "status": 200,
log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' "time": time
print(log) }
torch_gc() log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
return answer print(log)
torch_gc()
return answer
return await asyncio.to_thread(_sync_chat, history=history)
if __name__ == '__main__': if __name__ == '__main__':