Add history in API

API.py

@@ -1,33 +1,35 @@
-from typing import Optional
 from fastapi import FastAPI, Request
-from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModel
-import uvicorn, json, time, datetime, os, platform
+import uvicorn, json, datetime
 
 app = FastAPI()
+
+
 @app.post("/")
 async def create_item(request: Request):
-    global history, model, tokenizer
-    jsonPostRaw = await request.json()
-    jsonPost = json.dumps(jsonPostRaw)
-    jsonPostList = json.loads(jsonPost)
-    prompt = jsonPostList.get('prompt')
+    global model, tokenizer
+    json_post_raw = await request.json()
+    json_post = json.dumps(json_post_raw)
+    json_post_list = json.loads(json_post)
+    prompt = json_post_list.get('prompt')
+    history = json_post_list.get('history')
     response, history = model.chat(tokenizer, prompt, history=history)
     now = datetime.datetime.now()
     time = now.strftime("%Y-%m-%d %H:%M:%S")
     answer = {
-        "response":response,
-        "status":200,
-        "time":time
+        "response": response,
+        "history": history,
+        "status": 200,
+        "time": time
     }
-    log = "["+time+"] "+'device:"'+jsonPostList.get('device')+'", prompt:"'+prompt+'", response:"'+repr(response)+'"'
+    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
     print(log)
     return answer
 
+
 if __name__ == '__main__':
-    uvicorn.run('API:app',host='0.0.0.0',port=8000,workers=1)
+    uvicorn.run('API:app', host='0.0.0.0', port=8000, workers=1)
 
-history = []
 tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
-model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().quantize(4).cuda()
-model = model.eval()
+model = AutoModel.from_pretrained("THUDM/chatglm_6b", trust_remote_code=True).half().cuda()
+model.eval()