diff --git a/web_demo.py b/web_demo.py new file mode 100755 index 0000000..e5f15a1 --- /dev/null +++ b/web_demo.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 + +from transformers import AutoTokenizer +import gradio as gr +import mdtex2html +from transformers import AutoModelForCausalLM, AutoTokenizer +from transformers.generation import GenerationConfig +from argparse import ArgumentParser +import sys + +print("Call args:" + str(sys.argv)) +parser = ArgumentParser() +parser.add_argument("--share", action="store_true", default=False) +parser.add_argument("--inbrowser", action="store_true", default=False) +parser.add_argument("--server_port", type=int, default=80) +parser.add_argument("--server_name", type=str, default="0.0.0.0") +parser.add_argument("--exit", action="store_true", default=False) +parser.add_argument("--model_revision", type=str, default="") +args = parser.parse_args(sys.argv[1:]) +print("Args:" + str(args)) + +tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen-7B-Chat", trust_remote_code=True, resume_download=True +) + +model = AutoModelForCausalLM.from_pretrained( + "Qwen/Qwen-7B-Chat", + device_map="auto", + trust_remote_code=True, + resume_download=True, + **{"revision": args.model_revision} + if args.model_revision is not None + and args.model_revision != "" + and args.model_revision != "None" + else {}, +).eval() + +model.generation_config = GenerationConfig.from_pretrained( + "Qwen/Qwen-7B-Chat", trust_remote_code=True, resume_download=True +) + +if "exit" in args: + if args.exit: + sys.exit(0) + else: + del args.exit + +if "model_revision" in args: + del args.model_revision + + +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): + 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] = """") + gr.Markdown("""