From b61929e4d061b211bfdfb3ad128cc27c2ff910f5 Mon Sep 17 00:00:00 2001
From: eric <xiaoshuang.eric@gmail.com>
Date: Tue, 1 Aug 2023 22:10:53 +0800
Subject: [PATCH] update

---
 callback.py | 148 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 148 insertions(+)
 create mode 100644 callback.py

diff --git a/callback.py b/callback.py
new file mode 100644
index 0000000..9e50761
--- /dev/null
+++ b/callback.py
@@ -0,0 +1,148 @@
+from typing import Union
+from fastapi import FastAPI, Header, BackgroundTasks,Request,Body
+from urllib.parse import unquote
+import re
+import pickle
+import os
+from fastapi.responses import JSONResponse
+import asyncio
+from pydantic import BaseModel
+import time
+from transformers import AutoTokenizer, AutoModel
+import uvicorn, json, datetime
+import torch
+
+DEVICE = "cuda"
+DEVICE_ID = "0"
+CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE
+
+
+def torch_gc():
+    if torch.cuda.is_available():
+        with torch.cuda.device(CUDA_DEVICE):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+
+
+app = FastAPI()
+
+
+class Item(BaseModel):
+    query: Union[str, None] = None
+
+
+@app.get("/")
+def index():
+    return "hello"
+
+
+
+
+@app.get("/home")
+# def read_root():
+# return {"Hello": "World"}
+def read_items(
+    appid: Union[str, None] = Header(default=None),
+    bid: Union[str, None] = Header(default=None),
+    requestid: Union[str, None] = Header(default=None),
+    uid: Union[str, None] = Header(default=None),
+):
+    print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid)
+    return {"appid": appid, "bid": bid, "requestid": requestid, "uid": uid}
+
+
+@app.post("/glm")
+async def create_item(request: Request):
+    # 获取请求参数
+    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')  #默认空数组
+    max_length = json_post_list.get('max_length')
+    top_p = json_post_list.get('top_p')
+    temperature = json_post_list.get('temperature')
+
+    # 请求ai模型
+    response, history = 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)
+    now = datetime.datetime.now()
+    time = now.strftime("%Y-%m-%d %H:%M:%S")
+    answer = {
+        "response": response,
+        "history": history,
+        "status": 200,
+        "time": time
+    }
+    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
+    print(log)
+    torch_gc()
+    return answer
+
+
+# 通过网址发送闪信
+@app.post("/callback")
+async def callback(
+    # item: Item,
+    # inputArguments: Union[str, None] = None,
+    query=Body(None),
+    appid: Union[str, None] = Header(default=None),
+    bid: Union[str, None] = Header(default=None),
+    requestid: Union[str, None] = Header(default=None),
+    uid: Union[str, None] = Header(default=None),
+):
+    print('query',query)
+    print('query[query]',query["query"])
+
+    # 获取请求参数
+    global model, tokenizer
+    # time.sleep(10)
+    # print("item", item)
+    history = []  #默认空数组
+    prompt=''
+    try:
+        query = unquote(query["query"], "utf-8")
+        prompt = query #问题
+        
+    except:
+        prompt = ""
+   
+    max_length=None
+    top_p=None
+    temperature=None
+    # 请求ai模型
+    response, history = 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)
+    # log日志
+    now = datetime.datetime.now()
+    time = now.strftime("%Y-%m-%d %H:%M:%S")
+    log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
+    print(log)
+    torch_gc()
+    # return answer
+    content = {
+        "answer": response,
+        "answer_type": "text",
+    }
+    headers = {"Content-Type": "application/json"}
+    return JSONResponse(content=content, headers=headers)
+
+if __name__ == '__main__':
+    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
+    model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).cuda()
+    # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
+    # model_path = "THUDM/chatglm2-6b"
+    # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    # model = load_model_on_gpus(model_path, num_gpus=2)
+    model.eval()
+    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)
\ No newline at end of file