init project

README.md

@@ -1,13 +0,0 @@
----
-title: Chatglm2 6b Explorer
-emoji: 🐠
-colorFrom: green
-colorTo: blue
-sdk: gradio
-sdk_version: 3.37.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

apps/__init__.py

@@ -0,0 +1,3 @@
+from .translator import translator_demo
+from .simple_chat import simple_chat_demo
+from .instruction_chat import instruction_chat_demo

apps/components.py

@@ -0,0 +1,22 @@
+import gradio as gr
+
+
+def chat_accordion():
+    with gr.Accordion("参数设置", open=False):
+        temperature = gr.Slider(
+            minimum=0.1,
+            maximum=2.0,
+            value=0.8,
+            step=0.1,
+            interactive=True,
+            label="Temperature",
+        )
+        top_p = gr.Slider(
+            minimum=0.1,
+            maximum=0.99,
+            value=0.9,
+            step=0.01,
+            interactive=True,
+            label="top_p",
+        )
+    return temperature, top_p

apps/instruction_chat.py

@@ -0,0 +1,106 @@
+import traceback
+
+import gradio as gr
+
+from chatClient import ChatClient
+from apps.components import chat_accordion
+
+BOT_NAME = "ChatGLM2-6B"
+TITLE = """<h3 align="center">🤗 ChatGLM2-6B 预设指令对话</h3>"""
+RETRY_COMMAND = "/retry"
+DEFAULT_INSTRUCTIONS = "你是一个爱笑的机器人，名字叫小莫，在回答问题的时候会添加合适的emoji来表达情绪。"
+
+
+def chat(client: ChatClient):
+    with gr.Row():
+        with gr.Column(elem_id="chat_container", scale=3):
+            with gr.Row():
+                chatbot = gr.Chatbot(elem_id="chatbot")
+            with gr.Row():
+                inputs = gr.Textbox(
+                    placeholder=f"你好 {BOT_NAME} !",
+                    label="输入内容后点击回车",
+                    max_lines=3,
+                )
+            with gr.Row(elem_id="button_container"):
+                with gr.Column():
+                    retry_button = gr.Button("♻️ 重试上一轮对话")
+                with gr.Column():
+                    delete_turn_button = gr.Button("🧽 删除上一轮对话")
+                with gr.Column():
+                    clear_chat_button = gr.Button("✨ 删除全部对话历史")
+
+        with gr.Column(elem_id="param_container", scale=1):
+            with gr.Row():
+                with gr.Accordion("对话预设指令", open=True):
+                    instructions = gr.Textbox(
+                            placeholder="LLM instructions",
+                            value=DEFAULT_INSTRUCTIONS,
+                            lines=10,
+                            interactive=True,
+                            label="指令",
+                            max_lines=16,
+                            show_label=False,
+                        )
+            with gr.Row():
+                temperature, top_p = chat_accordion()
+
+    def run_chat(message: str, chat_history, instructions: str, temperature: float, top_p: float):
+        if not message or (message == RETRY_COMMAND and len(chat_history) == 0):
+            yield chat_history
+            return
+
+        if message == RETRY_COMMAND and chat_history:
+            prev_turn = chat_history.pop(-1)
+            user_message, _ = prev_turn
+            message = user_message
+
+        # chat_history = chat_history + [[message, ""]]
+        try:
+            stream = client.instruct_chat(
+                message,
+                chat_history,
+                instructions,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            for resp, history in stream:
+                chat_history = history
+                yield chat_history
+        except Exception as e:
+            if not chat_history:
+                chat_history = []
+            chat_history += [["有错误了", traceback.format_exc()]]
+            yield chat_history
+
+    def delete_last_turn(chat_history):
+        if chat_history:
+            chat_history.pop(-1)
+        return {chatbot: gr.update(value=chat_history)}
+
+    def run_retry(message: str, chat_history, instructions, temperature: float, top_p: float):
+        yield from run_chat(RETRY_COMMAND, chat_history, instructions, temperature, top_p)
+
+    def clear_chat():
+        return []
+
+    inputs.submit(
+        run_chat,
+        [inputs, chatbot, instructions, temperature, top_p],
+        outputs=[chatbot],
+        show_progress=False,
+    )
+    inputs.submit(lambda: "", inputs=None, outputs=inputs)
+    delete_turn_button.click(delete_last_turn, inputs=[chatbot], outputs=[chatbot])
+    retry_button.click(
+        run_retry,
+        [inputs, chatbot, instructions, temperature, top_p],
+        outputs=[chatbot],
+        show_progress=False,
+    )
+    clear_chat_button.click(clear_chat, [], chatbot)
+
+
+def instruction_chat_demo(client: ChatClient):
+    gr.HTML(TITLE)
+    chat(client)

apps/simple_chat.py

@@ -0,0 +1,90 @@
+import gradio as gr
+from chatClient import ChatClient
+import traceback
+from apps.components import chat_accordion
+
+BOT_NAME = "ChatGLM2-6B"
+TITLE = """<h3 align="center">🤖 通用对话</h3>"""
+RETRY_COMMAND = "/retry"
+
+
+def chat(client: ChatClient):
+    with gr.Row():
+        with gr.Column(elem_id="chat_container", scale=3):
+            with gr.Row():
+                chatbot = gr.Chatbot(elem_id="chatbot")
+            with gr.Row():
+                inputs = gr.Textbox(
+                    placeholder=f"你好 {BOT_NAME} !",
+                    label="输入内容后点击回车",
+                    max_lines=3,
+                )
+            with gr.Row(elem_id="button_container"):
+                with gr.Column():
+                    retry_button = gr.Button("♻️ 重试上一轮对话")
+                with gr.Column():
+                    delete_turn_button = gr.Button("🧽 删除上一轮对话")
+                with gr.Column():
+                    clear_chat_button = gr.Button("✨ 删除全部对话历史")
+
+        with gr.Column(elem_id="param_container", scale=1):
+            temperature, top_p = chat_accordion()
+
+    def run_chat(message: str, chat_history, temperature: float, top_p: float):
+        if not message or (message == RETRY_COMMAND and len(chat_history) == 0):
+            yield chat_history
+            return
+
+        if message == RETRY_COMMAND and chat_history:
+            prev_turn = chat_history.pop(-1)
+            user_message, _ = prev_turn
+            message = user_message
+
+        # chat_history = chat_history + [[message, ""]]
+        try:
+            stream = client.simple_chat(
+                message,
+                chat_history,
+                temperature=temperature,
+                top_p=top_p,
+            )
+            for resp, history in stream:
+                chat_history = history
+                yield chat_history
+        except Exception as e:
+            if not chat_history:
+                chat_history = []
+            chat_history += [["有错误了", traceback.format_exc()]]
+            yield chat_history
+
+    def delete_last_turn(chat_history):
+        if chat_history:
+            chat_history.pop(-1)
+        return {chatbot: gr.update(value=chat_history)}
+
+    def run_retry(message: str, chat_history, temperature: float, top_p: float):
+        yield from run_chat(RETRY_COMMAND, chat_history, temperature, top_p)
+
+    def clear_chat():
+        return []
+
+    inputs.submit(
+        run_chat,
+        [inputs, chatbot, temperature, top_p],
+        outputs=[chatbot],
+        show_progress=False,
+    )
+    inputs.submit(lambda: "", inputs=None, outputs=inputs)
+    delete_turn_button.click(delete_last_turn, inputs=[chatbot], outputs=[chatbot])
+    retry_button.click(
+        run_retry,
+        [inputs, chatbot, temperature, top_p],
+        outputs=[chatbot],
+        show_progress=False,
+    )
+    clear_chat_button.click(clear_chat, [], chatbot)
+
+
+def simple_chat_demo(client: ChatClient):
+    gr.HTML(TITLE)
+    chat(client)

apps/translator.py

@@ -0,0 +1,57 @@
+import traceback
+
+import gradio as gr
+
+from apps.components import chat_accordion
+
+IDEA_TITLE = "ChatGLM2-6B 翻译官"
+
+prompt_tmpl = """imagine you are a professional translator. Your task is translating the text around by ``` to Chinese.
+
+input text：
+
+```
+{input_text}
+```
+
+translation result:"""
+
+
+def translator_demo(client):
+
+    def stream_translate(input_text, temperature: float, top_p: float):
+        if not input_text:
+            return None
+        message = prompt_tmpl.format(input_text=input_text)
+        try:
+            stream = client.simple_chat(
+                message,
+                [],
+                temperature=temperature,
+                top_p=top_p,
+            )
+            for resp, _ in stream:
+                pass
+            return resp
+        except Exception as e:
+            return traceback.format_exc()
+    
+    def clear_content():
+        return None, None
+
+    with gr.Row():
+        with gr.Column():
+            inputs = gr.Textbox(label="请输入原文", max_lines=5)
+            gr.Dropdown(["en -> zh"], value="en -> zh", label="翻译语言")
+            temperature, top_p = chat_accordion()
+            with gr.Row(elem_id="button_container"):
+                with gr.Column():
+                    commit_btn = gr.Button(value="翻译", variant='primary')
+                with gr.Column():
+                    clear_btn = gr.Button(value="清空")
+
+        with gr.Column():
+            outputs = gr.Textbox(label="译文", max_lines=5)
+
+        commit_btn.click(stream_translate, inputs=[inputs, temperature, top_p], outputs=[outputs])
+        clear_btn.click(clear_content, inputs=None, outputs=[inputs, outputs])

chatClient.py

@@ -0,0 +1,99 @@
+import json
+
+from websockets.exceptions import ConnectionClosedOK
+from websockets.sync.client import connect
+
+from chatglm2_6b.modelClient import ChatGLM2
+import abc
+
+
+class ChatClient(abc.ABC):
+    @abc.abstractmethod
+    def simple_chat(self, query, history, temperature, top_p):
+        pass
+
+    @abc.abstractmethod
+    def instruct_chat(self, message, chat_history, instructions, temperature, top_p):
+        pass
+
+
+def format_chat_prompt(message: str, chat_history, instructions: str) -> str:
+    instructions = instructions.strip(" ").strip("\n")
+    prompt = f"对话背景设定:{instructions}"
+    for i, (user_message, bot_message) in enumerate(chat_history):
+        prompt = f"{prompt}\n\n[Round {i + 1}]\n\n问：{user_message}\n\n答：{bot_message}"
+    prompt = f"{prompt}\n\n[Round {len(chat_history)+1}]\n\n问：{message}\n\n答："
+    return prompt
+
+
+class ChatGLM2APIClient(ChatClient):
+    def __init__(self, ws_url=None):
+        self.ws_url = "ws://localhost:10001"
+        if ws_url:
+            self.ws_url = ws_url
+
+    def simple_chat(self, query, history, temperature, top_p):
+        """chatglm2-6b 模型定义的对话方法"""
+        url = f"{self.ws_url}/streamChat"
+        with connect(url) as websocket:
+            msg = json.dumps({
+                "query": query, "history": history, 
+                "temperature": temperature, "top_p": top_p,
+            })
+            websocket.send(msg)
+
+            data = None
+            try:
+                while True:
+                    data = websocket.recv()
+                    data = json.loads(data)
+                    yield data['resp'], data['history']
+            except ConnectionClosedOK:
+                print("generation is finished")
+
+    def instruct_chat(self, message, chat_history, instructions, temperature, top_p):
+        """基于chatglm2-6b text_generate 实现的基于预设指令的对话"""
+        url = f"{self.ws_url}/streamGenerate"
+
+        prompt = format_chat_prompt(message, chat_history, instructions)
+        chat_history = chat_history + [[message, ""]]
+        params = json.dumps({"prompt": prompt, "temperature": temperature, "top_p": top_p})
+        with connect(url) as websocket:
+            websocket.send(params)
+
+            data = None
+            try:
+                while True:
+                    data = websocket.recv()
+                    data = json.loads(data)
+                    resp = data['text']
+
+                    last_turn = list(chat_history.pop(-1))
+                    last_turn[-1] = resp
+                    chat_history = chat_history + [last_turn]
+                    yield resp, chat_history
+            except ConnectionClosedOK:
+                print("generation is finished")
+
+
+class ChatGLM2ModelClient(ChatClient):
+    def __init__(self, model_path=None):
+        self.model = ChatGLM2(model_path)
+
+    def simple_chat(self, query, history, temperature, top_p):
+        kwargs = {
+            "query": query, "history": history, 
+            "temperature": temperature, "top_p": top_p,
+        }
+        for resp, history in self.model.stream_chat(**kwargs):
+            yield resp, history
+
+    def instruct_chat(self, message, chat_history, instructions, temperature, top_p):
+        prompt = format_chat_prompt(message, chat_history, instructions)
+        chat_history = chat_history + [[message, ""]]
+        kwargs = {"prompt": prompt, "temperature": temperature, "top_p": top_p}
+        for resp in self.model.stream_generate(**kwargs):
+            last_turn = list(chat_history.pop(-1))
+            last_turn[-1] = resp
+            chat_history = chat_history + [last_turn]
+            yield resp, chat_history

chatglm2_6b/modelClient.py

@@ -0,0 +1,82 @@
+from typing import List, Tuple
+
+import torch
+from transformers import AutoTokenizer, AutoModel
+from transformers.generation.logits_process import LogitsProcessor
+from transformers.generation.utils import LogitsProcessorList
+
+DEFAULT_MODEL_PATH = "THUDM/chatglm2-6b"
+
+
+class InvalidScoreLogitsProcessor(LogitsProcessor):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+        if torch.isnan(scores).any() or torch.isinf(scores).any():
+            scores.zero_()
+            scores[..., 5] = 5e4
+        return scores
+
+
+class ChatGLM2(object):
+    def __init__(self, model_path=None):
+        if not model_path:
+            self.model_path = DEFAULT_MODEL_PATH
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
+        model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).half().cuda()
+        self.model = model.eval()
+
+    def generate(
+        self,
+        prompt: str,
+        do_sample: bool = True,
+        max_length: int = 8192,
+        num_beams: int = 1,
+        temperature: float = 0.8,
+        top_p: float = 0.8,
+    ):
+        logits_processor = LogitsProcessorList()
+        logits_processor.append(InvalidScoreLogitsProcessor())
+        gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "do_sample": do_sample, "top_p": top_p,
+                      "temperature": temperature, "logits_processor": logits_processor}
+        inputs = self.tokenizer([prompt], return_tensors="pt")
+        inputs = inputs.to(self.model.device)
+        outputs = self.model.generate(**inputs, **gen_kwargs)
+        outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
+        response = self.tokenizer.decode(outputs)
+        response = self.model.process_response(response)
+        return response
+
+    def stream_generate(
+        self,
+        prompt: str,
+        do_sample: bool = True,
+        max_length: int = 8192,
+        temperature: float = 0.8,
+        top_p: float = 0.8,
+    ):
+        logits_processor = LogitsProcessorList()
+        logits_processor.append(InvalidScoreLogitsProcessor())
+        gen_kwargs = {"max_length": max_length, "do_sample": do_sample, "top_p": top_p,
+                      "temperature": temperature, "logits_processor": logits_processor}
+        inputs = self.tokenizer([prompt], return_tensors="pt")
+        inputs = inputs.to(self.model.device)
+        for outputs in self.model.stream_generate(**inputs, **gen_kwargs):
+            outputs = outputs.tolist()[0][len(inputs["input_ids"][0]):]
+            response = self.tokenizer.decode(outputs)
+            if response and response[-1] != "�":
+                response = self.model.process_response(response)
+                yield response
+
+    def stream_chat(
+        self, 
+        query: str, 
+        history: List[Tuple[str, str]],
+        max_length: int = 8192, 
+        do_sample=True,
+        top_p=0.8,
+        temperature=0.8
+    ):
+        stream = self.model.stream_chat(self.tokenizer, query, history, 
+            max_length=max_length, do_sample=do_sample, top_p=top_p, temperature=temperature)
+        for resp, new_history in stream:
+            yield resp, new_history
+    

chatglm2_6b/server.py

@@ -0,0 +1,49 @@
+import functools
+
+import anyio
+from fastapi import FastAPI, WebSocket
+from pydantic import BaseModel
+
+from chatglm2_6b.modelClient import ChatGLM2
+from config import Settings
+
+app = FastAPI()
+
+chat_glm2 = ChatGLM2(Settings.CHATGLM_MODEL_PATH)
+
+
+class ChatParams(BaseModel):
+    prompt: str
+    do_sample: bool = True
+    max_length: int = 2048
+    temperature: float = 0.8
+    top_p: float = 0.8
+
+
+@app.post("/generate")
+def generate(params: ChatParams):
+    input_params = params.dict()
+    text = chat_glm2.generate(**input_params)
+    return {"text": text}
+
+
+@app.websocket("/streamGenerate")
+async def stream_generate(websocket: WebSocket):
+    await websocket.accept()
+    params = await websocket.receive_json()
+    func = functools.partial(chat_glm2.stream_generate, **params)
+    stream = await anyio.to_thread.run_sync(func)
+    for resp in stream:
+        await websocket.send_json({"text": resp})
+    await websocket.close()
+
+
+@app.websocket("/streamChat")
+async def stream_chat(websocket: WebSocket):
+    await websocket.accept()
+    params = await websocket.receive_json()
+    func = functools.partial(chat_glm2.stream_chat, **params)
+    stream = await anyio.to_thread.run_sync(func)
+    for resp, history in stream:
+        await websocket.send_json({"resp": resp, "history": history})
+    await websocket.close()

config.py

@@ -0,0 +1,7 @@
+import os
+
+
+class Settings:
+    CHAT_CLIENT = os.environ.get('CHAT_CLIENT', "ChatGLM2APIClient")
+    MODEL_WS_URL = os.environ.get('MODEL_WS_URL', "ws://localhost:10001")
+    CHATGLM_MODEL_PATH = os.environ.get('CHATGLM_MODEL_PATH', "THUDM/chatglm2-6b")

gallery.gradio.py

@@ -0,0 +1,47 @@
+import gradio as gr
+from apps import translator_demo, simple_chat_demo, instruction_chat_demo
+from chatClient import ChatClient, ChatGLM2APIClient, ChatGLM2ModelClient
+from config import Settings
+
+TITLE = """<h2 align="center">🚀 ChatGLM2-6B apps gallery</h2>"""
+
+demo_register = {
+    "通用对话": simple_chat_demo,
+    "预设指令对话": instruction_chat_demo,
+    "翻译器": translator_demo,
+}
+
+
+def get_gallery(client: ChatClient):
+    with gr.Blocks(
+        # css=None
+        # css="""#chat_container {width: 700px; margin-left: auto; margin-right: auto;}
+        #        #button_container {width: 700px; margin-left: auto; margin-right: auto;}
+        #        #param_container {width: 700px; margin-left: auto; margin-right: auto;}"""
+        css="""#chatbot {
+            font-size: 14px;
+            min-height: 300px;
+        }"""
+    ) as demo:
+        gr.HTML(TITLE)
+        for name, demo_func in demo_register.items():
+            with gr.Tab(name):
+                demo_func(client)
+    return demo
+
+
+def build_client():
+    client_class = Settings.CHAT_CLIENT
+    if client_class == 'ChatGLM2ModelClient':
+
+        return ChatGLM2ModelClient(Settings.CHATGLM_MODEL_PATH)
+    if client_class == 'ChatGLM2APIClient':
+        return ChatGLM2APIClient(Settings.MODEL_WS_URL)
+    raise Exception(f"Wrong ChatClient: {client_class}")
+
+
+if __name__ == "__main__":
+    client = build_client()
+    demo = get_gallery(client)
+    demo.queue(max_size=128, concurrency_count=16)
+    demo.launch()

runserver.py

@@ -0,0 +1,11 @@
+import uvicorn
+
+from chatglm2_6b.server import app
+
+
+def runserver():
+    uvicorn.run(app, host="0.0.0.0", port=10001)
+
+
+if __name__ == '__main__':
+    runserver()