diff --git a/run_web_demo.sh b/run_web_demo.sh new file mode 100755 index 0000000..f3f9a8b --- /dev/null +++ b/run_web_demo.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash + +cd "$(dirname "$0")" +thisDir=$(pwd) + +function performInstall() { + set -e + + pushd "$thisDir" + pip3 install -r requirements.txt + pip3 install gradio mdtex2html scipy + + if [[ ! -d flash-attention ]]; then + if ! git clone -b v1.0.8 https://github.com/Dao-AILab/flash-attention; then + echo "Clone flash-attention failed, please install it manually." + return 0 + fi + fi + + cd flash-attention && + pip3 install . && + pip3 install csrc/layer_norm && + pip3 install csrc/rotary || + echo "Install flash-attention failed, please install it manually." + popd +} + +echo "Starting WebUI..." + +if ! python3 web_demo.py; then + echo "Run demo failed, install the deps and try again? (y/n)" + # auto perform install if in docker + if [[ -t 0 ]] && [[ -t 1 ]] && [[ ! -f "/.dockerenv" ]]; then + read doInstall + else + doInstall="y" + fi + + if ! [[ "$doInstall" =~ y|Y ]]; then + exit 1 + fi + + echo "Installing deps, and try again..." + performInstall && python3 web_demo.py +fi diff --git a/web_demo.py b/web_demo.py new file mode 100755 index 0000000..8b37b02 --- /dev/null +++ b/web_demo.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 + +""" Ref: https://github.com/THUDM/ChatGLM2-6B/blob/main/web_demo.py """ + +from transformers import AutoTokenizer +import gradio as gr +import mdtex2html +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.generation import GenerationConfig +import sys + +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True) +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat", device_map="auto", trust_remote_code=True).eval() +model.generation_config = GenerationConfig.from_pretrained("Qwen/Qwen-7B-Chat", trust_remote_code=True) + +if len(sys.argv) > 1 and sys.argv[1] == "--exit": + exit(0) + +def postprocess(self, y): + if y is None: + return [] + for i, (message, response) in enumerate(y): + y[i] = ( + None if message is None else mdtex2html.convert((message)), + None if response is None else mdtex2html.convert(response), + ) + return y + + +gr.Chatbot.postprocess = postprocess + + +def parse_text(text): + """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/""" + lines = text.split("\n") + lines = [line for line in lines if line != ""] + count = 0 + for i, line in enumerate(lines): + if "```" in line: + count += 1 + items = line.split('`') + if count % 2 == 1: + lines[i] = f'
'
+ else:
+ lines[i] = f'
'
+ else:
+ if i > 0:
+ if count % 2 == 1:
+ line = line.replace("`", "\`")
+ line = line.replace("<", "<")
+ line = line.replace(">", ">")
+ line = line.replace(" ", " ")
+ line = line.replace("*", "*")
+ line = line.replace("_", "_")
+ line = line.replace("-", "-")
+ line = line.replace(".", ".")
+ line = line.replace("!", "!")
+ line = line.replace("(", "(")
+ line = line.replace(")", ")")
+ line = line.replace("$", "$")
+ lines[i] = "