diff --git a/cli_demo_vision.py b/cli_demo_vision.py new file mode 100644 index 0000000..f0e217f --- /dev/null +++ b/cli_demo_vision.py @@ -0,0 +1,65 @@ +import os +import platform +import signal +from transformers import AutoTokenizer, AutoModel +import readline + +tokenizer = AutoTokenizer.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True).half().cuda() +model = model.eval() + +os_name = platform.system() +clear_command = 'cls' if os_name == 'Windows' else 'clear' +stop_stream = False + + +def build_prompt(history, prefix): + prompt = prefix + for query, response in history: + prompt += f"\n\n用户:{query}" + prompt += f"\n\nChatGLM-6B:{response}" + return prompt + + +def signal_handler(signal, frame): + global stop_stream + stop_stream = True + + +def main(): + global stop_stream + while True: + history = [] + prefix = "欢迎使用 VisualGLM-6B 模型,输入图片路径和内容即可进行对话,clear 清空对话历史,stop 终止程序" + print(prefix) + image_path = input("\n请输入图片路径:") + if image_path == "stop": + break + prefix = prefix + "\n" + image_path + query = "描述这张图片。" + while True: + count = 0 + for response, history in model.stream_chat(tokenizer, image_path, query, history=history): + if stop_stream: + stop_stream = False + break + else: + count += 1 + if count % 8 == 0: + os.system(clear_command) + print(build_prompt(history, prefix), flush=True) + signal.signal(signal.SIGINT, signal_handler) + os.system(clear_command) + print(build_prompt(history, prefix), flush=True) + query = input("\n用户:") + if query.strip() == "stop": + break + if query.strip() == "clear": + history = [] + os.system(clear_command) + print(prefix) + continue + + +if __name__ == "__main__": + main() diff --git a/web_demo_vision.py b/web_demo_vision.py new file mode 100644 index 0000000..79f6b52 --- /dev/null +++ b/web_demo_vision.py @@ -0,0 +1,120 @@ +from transformers import AutoModel, AutoTokenizer +import gradio as gr +import mdtex2html + +tokenizer = AutoTokenizer.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True) +model = AutoModel.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True).half().cuda() +model = model.eval() + +"""Override Chatbot.postprocess""" + + +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] = "