diff --git a/README.md b/README.md index 554a747..7e3c7e5 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,17 @@ ChatGLM-6B 使用了和 ChatGPT 相似的技术,针对中文问答和对话进 **[2023/03/19]** 增加流式输出接口 `stream_chat`,已更新到网页版和命令行 Demo。修复输出中的中文标点。增加量化后的模型 [ChatGLM-6B-INT4](https://huggingface.co/THUDM/chatglm-6b-int4) + +## 友情链接 +以下是部分基于本仓库开发的开源项目: +* [ChatGLM-MNN](https://github.com/wangzhaode/ChatGLM-MNN): 一个基于 MNN 的 ChatGLM-6B C++ 推理实现,支持根据显存大小自动分配计算任务给 GPU 和 CPU +* [ChatGLM-Tuning](https://github.com/mymusise/ChatGLM-Tuning): 基于 LoRA 对 ChatGLM-6B 进行微调 + +以下是部分针对本项目的教程/文档: +* [Windows部署文档](https://github.com/ZhangErling/ChatGLM-6B/blob/main/deployment_windows.md) + +如果你有其他好的项目/教程的话,欢迎参照上述格式添加到README中并提出 [PR](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork). + ## 使用方式 ### 硬件需求 diff --git a/README_en.md b/README_en.md index 3a8b17b..7b84f9c 100644 --- a/README_en.md +++ b/README_en.md @@ -13,6 +13,13 @@ Try the [online demo](https://huggingface.co/spaces/ysharma/ChatGLM-6b_Gradio_St **[2023/03/19]** Add streaming output function `stream_chat`, already applied in web and CLI demo. Fix Chinese punctuations in output. Add quantized model [ChatGLM-6B-INT4](https://huggingface.co/THUDM/chatglm-6b-int4). +## Projects +The following are some open source projects developed based on this repository: +* [ChatGLM-MNN](https://github.com/wangzhaode/ChatGLM-MNN): An [MNN](https://github.com/alibaba/MNN)-based implementation of ChatGLM-6B C++ inference, which supports automatic allocation of computing tasks to GPU and CPU according to the size of GPU memory +* [ChatGLM-Tuning](https://github.com/mymusise/ChatGLM-Tuning): Fine-tuning ChatGLM-6B based on LoRA + +If you have other good projects, please refer to the above format to add to README and propose [PR](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request-from-a-fork). + ## Getting Started ### Hardware Requirements diff --git a/api.py b/api.py index 10f70b6..693c70a 100644 --- a/api.py +++ b/api.py @@ -1,6 +1,19 @@ from fastapi import FastAPI, Request 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() @@ -13,7 +26,15 @@ async def create_item(request: Request): 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) + max_length = json_post_list.get('max_length') + top_p = json_post_list.get('top_p') + temperature = json_post_list.get('temperature') + 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 = { @@ -24,12 +45,12 @@ async def create_item(request: Request): } log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"' print(log) + torch_gc() return answer if __name__ == '__main__': - uvicorn.run('api:app', host='0.0.0.0', port=8000, workers=1) - -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() + 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() + uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) diff --git a/cli_demo.py b/cli_demo.py index 8a043fb..1c3ff2b 100644 --- a/cli_demo.py +++ b/cli_demo.py @@ -1,5 +1,6 @@ import os import platform +import signal from transformers import AutoTokenizer, AutoModel tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True) @@ -8,6 +9,7 @@ model = model.eval() os_name = platform.system() clear_command = 'cls' if os_name == 'Windows' else 'clear' +stop_stream = False def build_prompt(history): @@ -18,8 +20,14 @@ def build_prompt(history): return prompt +def signal_handler(signal, frame): + global stop_stream + stop_stream = True + + def main(): history = [] + global stop_stream print("欢迎使用 ChatGLM-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序") while True: query = input("\n用户:") @@ -32,10 +40,15 @@ def main(): continue count = 0 for response, history in model.stream_chat(tokenizer, query, history=history): - count += 1 - if count % 8 == 0: - os.system(clear_command) - print(build_prompt(history), flush=True) + if stop_stream: + stop_stream = False + break + else: + count += 1 + if count % 8 == 0: + os.system(clear_command) + print(build_prompt(history), flush=True) + signal.signal(signal.SIGINT, signal_handler) os.system(clear_command) print(build_prompt(history), flush=True)