feat: chatglm stream api

cli_chatglm_with_api.py

@@ -0,0 +1,157 @@
+#!/usr/bin/python3
+import os
+import platform
+import requests
+import json
+import time
+
+os_name = platform.system()
+clear_command = "cls" if os_name == "Windows" else "clear"
+stream = True
+api_url = os.environ.get("ChatGLM_API", "http://127.0.0.1:8888") + "/chat"
+max_history_len = 6
+max_role_history_len = 6
+welcome_text = "欢迎使用 ChatGLM-6B 模型，输入内容即可进行对话，role 设定角色，clear 清空对话历史，stop 终止程序"
+failed_resp_test = "AI 我呀，有点想不通了..."
+secret = os.environ.get("ChatGLM_SECRET", "721d95ac31da59fa022ec8c12f72f597")
+type_wait_time = 0.05
+
+headers = {
+    "Content-Type": "application/json",
+    "Authorization": secret,
+}
+
+
+def request_chat(data):
+    response = requests.post(api_url, headers=headers, json=data)
+    res = None
+    if response.status_code == 200:
+        res = response.json()
+    else:
+        print(
+            "Request failed with code {}.".format(response.status_code),
+            end="\n\n",
+        )
+    return res
+
+
+def request_stream_chat(data):
+    response = requests.post(api_url, headers=headers, json=data, stream=True)
+    if response.status_code == 200:
+        for lines in response.iter_lines(decode_unicode=True):
+            if lines:
+                data = json.loads(lines)
+                if data.get("status") == 200:
+                    if data.get("stop", True):
+                        break
+                    yield data
+                else:
+                    print(
+                        f"\n\nSystem Error: Status {data.get('status')}",
+                        end="\n\n",
+                    )
+                    break
+    else:
+        print(f"\n\nSystem Error: Status {response.status_code}")
+
+
+def main():
+    global stream
+    role_history = []
+    history = []
+    print(welcome_text, end="\n\n")
+    while True:
+        query = input("\n用户：")
+        print("", end="\n\n")
+
+        if query.strip() == "stop":
+            break
+        if query.strip() == "clear":
+            history = []
+            role_history = []
+            os.system(clear_command)
+            print(welcome_text, end="\n\n")
+            continue
+        if query.strip() == "role":
+            print(
+                f"请输入角色设定（注意提问方式），exit 取消设定, ok 完成设定，不超过{max_role_history_len}",
+                end="\n\n",
+            )
+            cancled = False
+            new_role = []
+            for i in range(1, max_role_history_len + 1):
+                query = input("\n设定 " + str(i) + "：")
+                if query.strip() == "exit":
+                    cancled = True
+                    break
+                if query.strip() == "ok":
+                    break
+                req_data = {
+                    "prompt": query,
+                    "history": new_role,
+                    "stream": False,
+                }
+
+                print("\n\nChatGLM-6B 记录中......", end="\n\n")
+
+                if stream:
+                    req_data["stream"] = True
+                    response = None
+                    for res_data in request_stream_chat(req_data):
+                        if res_data:
+                            response = res_data
+                else:
+                    response = request_chat(req_data)
+                if response:
+                    new_role = response.get("history", history)
+                else:
+                    print("该设定失败！", end="\n\n")
+            if not cancled:
+                role_history = new_role
+                history = role_history
+                print("设定角色成功！", end="\n\n")
+                print(history, end="\n\n")
+            continue
+
+        if len(role_history) > 0:
+            if len(history) > max_history_len + len(role_history):
+                history = role_history + history[-max_history_len:]
+        else:
+            history = history[-max_history_len:]
+
+        print("ChatGLM-6B：", end="")
+
+        # stream chat
+        if stream:
+            completed = None
+            last_stop = 0
+            for res_data in request_stream_chat(
+                {"prompt": query, "history": history, "stream": True}
+            ):
+                if res_data:
+                    text = res_data.get("response", failed_resp_test)
+                    print(text[last_stop:], end="", flush=True)
+                    last_stop = len(text)
+                    completed = res_data
+            print("", end="\n\n", flush=True)
+
+            if completed:
+                history = completed.get("history", history)
+
+        else:
+            res_data = request_chat(
+                {"prompt": query, "history": history, "stream": False}
+            )
+            if res_data:
+                history = res_data.get("history", history)
+                text = res_data.get("response", failed_resp_test)
+                for i in range(0, len(text), 8):
+                    print(text[i : i + 8], end="", flush=True)
+                    time.sleep(type_wait_time)
+                print("", end="\n\n", flush=True)
+            else:
+                print(failed_resp_test, end="\n\n")
+
+
+if __name__ == "__main__":
+    main()

stream_flask_api.py

@@ -0,0 +1,152 @@
+import os
+
+from flask import Flask, Response, request, redirect
+from transformers import AutoTokenizer, AutoModel
+import json, time, datetime, hashlib, redis
+import torch
+
+DEVICE = "cuda"
+DEVICE_ID = "0"
+CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE
+GLM_SECRET_KEY = os.environ.get(
+    "GLM_SECRET_KEY", "721d95ac31da59fa022ec8c12f72f597"
+)
+
+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()
+
+app = Flask(__name__)
+mq = redis.StrictRedis()
+
+
+def torch_gc():
+    if torch.cuda.is_available():
+        with torch.cuda.device(CUDA_DEVICE):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+
+def wrap_answer(prompt, response, history, stop=True):
+    now = datetime.datetime.now()
+    time = now.strftime("%Y-%m-%d %H:%M:%S")
+    answer = {
+        "response": response,
+        "history": history,
+        "stop": stop,
+        "status": 200,
+        "time": time,
+    }
+    log = (
+        "["
+        + time
+        + "] "
+        + '", prompt:"'
+        + prompt
+        + '", response:"'
+        + repr(response)
+        + '"'
+    )
+    print(log)
+    return answer
+
+
+def chat(prompt, history=[], max_length=None, top_p=None, temperature=None):
+    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,
+    )
+
+    return wrap_answer(prompt, response, history)
+
+
+def stream_chat(
+    prompt, history=[], max_length=None, top_p=None, temperature=None
+):
+    global model, tokenizer
+    for response, history in model.stream_chat(
+        tokenizer,
+        prompt,
+        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,
+    ):
+        answer = wrap_answer(prompt, response, history, stop=False)
+        yield answer
+
+    yield {"stop": True, "status": 200}
+
+
+@app.post("/chat")
+async def chat_api():
+    secret = request.headers.get("Authorization")
+    if secret != GLM_SECRET_KEY:
+        return {"status": 403}
+
+    global model, tokenizer
+    json_post_raw = request.json
+    json_post = json.dumps(json_post_raw)
+    json_post_list = json.loads(json_post)
+    stream = json_post_list.get("stream")
+
+    if stream:
+        key = hashlib.md5(
+            (f"{secret}-{int(time.time() * 1000)}").encode()
+        ).hexdigest()
+        mq.set(f"prompt-{key}", json_post)
+
+        return redirect(f"/stream/{key}", code=307)
+    else:
+        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")
+
+        answer = chat(prompt, history, max_length, top_p, temperature)
+    torch_gc()
+
+    return answer
+
+
+@app.route("/stream/<key>", methods=["GET", "POST"])
+def sse(key):
+    secret = request.headers.get("Authorization")
+    if secret != GLM_SECRET_KEY:
+        return {"status": 403}
+
+    if key:
+        post_data = mq.get(f"prompt-{key}")
+        if post_data:
+            json_post_list = json.loads(post_data)
+            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")
+
+            def generate():
+                for answer in stream_chat(
+                    prompt, history, max_length, top_p, temperature
+                ):
+                    yield "\n" + json.dumps(answer) + "\n"
+
+            return Response(generate(), mimetype="text/event-stream")
+
+    return {"status": 200}
+
+
+if __name__ == "__main__":
+    # pass
+    app.run(threaded=True, host="0.0.0.0", port=8888)