Init commit

LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023, OedoSoldier
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,13 +1,19 @@
----
-title: Chatglm Webui
-emoji: 📊
-colorFrom: yellow
-colorTo: indigo
-sdk: gradio
-sdk_version: 3.21.0
-app_file: app.py
-pinned: false
-license: bsd-3-clause
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ChatGLM WebUI
+
+基于 [ChatGPT WebUI](https://github.com/dotmet/chatgpt_webui) 的一个简单的 [ChatGLM](https://github.com/THUDM/ChatGLM-6B) WebUI。
+
+## 环境安装
+
+`pip install -r requirements.txt`
+
+Gradio 版本必须大于 3.21.0！
+
+Windows 下 CPU 运行需要安装编译器，可参考 (https://www.freesion.com/article/4185669814/)。
+
+## 运行
+
+通过 `git clone https://huggingface.co/THUDM/chatglm-6b-int4` 下载模型文件到根目录下然后 `python main.py` 即可，默认状态至少需要 4 GB 显存（CPU 运行则需要 5.2 GB 内存）。
+
+### 实参
+
+ - `--path`：指定模型所在文件夹

app.py

@@ -0,0 +1,207 @@
+import argparse
+import os
+import json
+import datetime
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description='ChatGLM Arguments')
+
+    parser.add_argument('--path', default='THUDM/chatglm-6b-int4', help='The path of ChatGLM model')
+
+    return parser.parse_args()
+
+
+args = get_args()
+
+if not os.path.isdir(args.path):
+    raise FileNotFoundError('Model not found')
+
+if torch.cuda.is_available():
+    device = 'cuda'
+else:
+    device = 'cpu'
+
+tokenizer = AutoTokenizer.from_pretrained(args.path, trust_remote_code=True)
+
+if device == 'cuda':
+    model = AutoModel.from_pretrained(args.path, trust_remote_code=True).half().cuda()
+else:
+    model = AutoModel.from_pretrained(args.path, trust_remote_code=True).float()
+
+model = model.eval()
+
+
+def parse_text(text):
+    lines = text.split('\n')
+    for i, line in enumerate(lines):
+        if '```' in line:
+            item = line.split('`')[-1]
+            if item:
+                lines[i] = f'<pre><code class="{item}">'
+            else:
+                lines[i] = '</code></pre>'
+        else:
+            if i > 0:
+                line = line.replace('<', '&lt;').replace('>', '&gt;')
+                lines[i] = f'<br/>{line}'
+    return ''.join(lines)
+
+
+def chat_wrapper(query, styled_history, history, max_length, top_p, temperature, memory_limit):
+    if query == '':
+        return [], [], '', *gr_hide()
+    if memory_limit == 0:
+        history = []
+        styled_history = []
+    elif memory_limit > 0:
+        history = history[-memory_limit:]
+        styled_history = styled_history[-memory_limit:]
+    flag = True
+    styled_history_pos = 0
+    for message, history in model.stream_chat(tokenizer, query, history=history,
+                                  max_length=max_length, top_p=top_p, temperature=temperature):
+        if flag:
+            styled_history.append((parse_text(query), parse_text(message)))
+            styled_history_pos = len(styled_history) - 1
+            flag = False
+        else:
+            styled_history[styled_history_pos] = (parse_text(query), parse_text(message))
+        yield styled_history, history, '', *gr_hide()
+
+
+def regenerate_wrapper(styled_history, history, max_length, top_p, temperature, memory_limit):
+    if not history:
+        return [], [], '', *gr_hide()
+
+    styled_history, history, query, _, _, _ = edit_wrapper(styled_history, history)
+    for ret in chat_wrapper(query, styled_history, history, max_length, top_p, temperature, memory_limit):
+        yield ret
+
+
+def edit_wrapper(styled_history, history):
+    if len(history) == 0:
+        return [], [], ''
+    query = history[-1][0]
+    history = history[:-1]
+    styled_history = styled_history[:-1]
+    return styled_history, history, query, *gr_hide()
+
+
+def reset_history():
+    return [], [], '', *gr_hide()
+
+
+def save_history(history):
+    os.makedirs('log', exist_ok=True)
+    dict_list = [{'input': q, 'output': a} for q, a in history]
+    with open(f'log/{datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")}.json', 'w', encoding='utf-8') as f:
+        json.dump(dict_list, f, ensure_ascii=False, indent=2)
+
+
+def save_config(max_length=2048, top_p=0.7, temperature=0.95, memory_limit=-1.0):
+    with open('config.json', 'w') as f:
+        json.dump({'max_length': max_length, 'top_p': top_p, 'temperature': temperature, 'memory_limit': memory_limit}, f, indent=2)
+
+
+def load_history(file, styled_history, history):
+    current_styled_history, current_history = styled_history.copy(), history.copy()
+    try:
+        with open(file.name, 'r', encoding='utf-8') as f:
+            dict_list = json.load(f)
+        history = [(item['input'], item['output']) for item in dict_list]
+        styled_history = [(parse_text(item['input']), parse_text(item['output'])) for item in dict_list]
+    except BaseException:
+        return current_styled_history, current_history, ''
+    return styled_history, history, '', *gr_hide()
+
+
+def gr_show_and_load(history, evt: gr.SelectData):
+    if evt.index[1] == 0:
+        label = f'Editing Q{evt.index[0]}：'
+    else:
+        label = f'Editing Q{evt.index[0]}：'
+    return {'visible': True, '__type__': 'update'}, {'value': history[evt.index[0]][evt.index[1]], 'label': label, '__type__': 'update'}, evt.index
+
+
+def update_history(styled_history, history, log, idx):
+    if log == '':
+        return styled_history, history, {'visible': True, '__type__': 'update'},  {'value': history[idx[0]][idx[1]], '__type__': 'update'}, idx
+
+    def swap_value(lst, idx, value):
+        lst[idx[0]] = tuple(value if j == idx[1] else elem for j, elem in enumerate(lst[idx[0]]))
+        return lst
+    styled_history = swap_value(styled_history, idx, parse_text(log))
+    history = swap_value(history, idx, log)
+    return styled_history, history, *gr_hide()
+
+
+def gr_hide():
+    return {'visible': False, '__type__': 'update'}, {'value': '', 'label': '', '__type__': 'update'}, []
+
+
+with gr.Blocks() as demo:
+    if not os.path.isfile('config.json'):
+        save_config()
+
+    with open('config.json', 'r', encoding='utf-8') as f:
+        configs = json.loads(f.read())
+
+    gr.Markdown('''<h1><center>ChatGLM Demo</center></h1>''')
+
+    with gr.Row():
+        max_length = gr.Slider(minimum=4.0, maximum=4096.0, step=4.0, label='Max Length', value=configs['max_length'])
+        top_p = gr.Slider(minimum=0.01, maximum=1.0, step=0.01, label='Top P', value=configs['top_p'])
+        temperature = gr.Slider(minimum=0.01, maximum=2.0, step=0.01, label='Temperature', value=configs['temperature'])
+        memory_limit = gr.Slider(minimum=-1.0, maximum=20.0, step=1.0, label='Memory Limit', value=configs['memory_limit'])
+    save_conf = gr.Button('保存设置', visible=False)
+
+    gr.Markdown('''<h2>Hint: click on a chat bubble to edit chat history</h2>''')
+
+    state = gr.State([])
+    chatbot = gr.Chatbot(elem_id='chatbot', show_label=False)
+    with gr.Row(visible=False) as edit_log:
+        with gr.Column():
+            log = gr.Textbox()
+            with gr.Row():
+                submit_log = gr.Button('保存')
+                cancel_log = gr.Button('取消')
+    log_idx = gr.State([])
+
+    message = gr.Textbox(placeholder='Input your message', label='Q:')
+
+    with gr.Row():
+        submit = gr.Button('Submit')
+        edit = gr.Button('Edit last question')
+        regen = gr.Button('Re-generate')
+
+    delete = gr.Button('Reset chat')
+
+    with gr.Row():
+        save = gr.Button('保存对话（在 `log` 文件夹下）', visible=False)
+        load = gr.UploadButton('读取对话', file_types=['file'], file_count='single', visible=False)
+
+    input_list = [message, chatbot, state, max_length, top_p, temperature, memory_limit]
+    output_list = [chatbot, state, message]
+    edit_list = [edit_log, log, log_idx]
+
+    save_conf.click(save_config, inputs=input_list[3:])
+    load.upload(load_history, inputs=[load, chatbot, state], outputs=output_list + edit_list)
+    save.click(save_history, inputs=[state])
+    message.submit(chat_wrapper, inputs=input_list, outputs=output_list + edit_list)
+    submit.click(chat_wrapper, inputs=input_list, outputs=output_list + edit_list)
+    edit.click(edit_wrapper, inputs=input_list[1:3], outputs=output_list + edit_list)
+    regen.click(regenerate_wrapper, inputs=input_list[1:], outputs=output_list + edit_list)
+    delete.click(reset_history, outputs=output_list + edit_list)
+    chatbot.select(gr_show_and_load, inputs=[state], outputs=edit_list)
+    edit_kwargs = {'inputs': [chatbot, state, log, log_idx], 'outputs': [chatbot, state] + edit_list} 
+    log.submit(update_history, **edit_kwargs)
+    submit_log.click(update_history, **edit_kwargs)
+    cancel_log.click(gr_hide, outputs=edit_list)
+
+
+if __name__ == '__main__':
+    demo.queue(concurrency_count=5, max_size=20).launch(debug=True)

config.json

@@ -0,0 +1,6 @@
+{
+  "max_length": 2048,
+  "top_p": 0.7,
+  "temperature": 0.95,
+  "memory_limit": -1
+}

requirements.txt

@@ -0,0 +1,6 @@
+protobuf>=3.19.5,<3.20.1
+transformers>=4.26.1
+icetk
+cpm_kernels
+torch>=1.10
+gradio>=3.21.0