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] = "
"+line + text = "".join(lines) + return text + + +def predict(input, image_path, chatbot, max_length, top_p, temperature, history): + if image_path is None: + return [(input, "图片为空!请重新上传图片并重试。")] + chatbot.append((parse_text(input), "")) + for response, history in model.stream_chat(tokenizer, image_path, input, history, max_length=max_length, top_p=top_p, + temperature=temperature): + chatbot[-1] = (parse_text(input), parse_text(response)) + + yield chatbot, history + + +def predict_new_image(image_path, chatbot, max_length, top_p, temperature): + input, history = "描述这张图片。", [] + chatbot.append((parse_text(input), "")) + for response, history in model.stream_chat(tokenizer, image_path, input, history, max_length=max_length, + top_p=top_p, + temperature=temperature): + chatbot[-1] = (parse_text(input), parse_text(response)) + + yield chatbot, history + + +def reset_user_input(): + return gr.update(value='') + + +def reset_state(): + return None, [], [] + + +with gr.Blocks() as demo: + gr.HTML("""

VisualGLM

""") + + image_path = gr.Image(type="filepath", label="Image Prompt", value=None) + chatbot = gr.Chatbot() + with gr.Row(): + with gr.Column(scale=4): + with gr.Column(scale=12): + user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style( + container=False) + with gr.Column(min_width=32, scale=1): + submitBtn = gr.Button("Submit", variant="primary") + with gr.Column(scale=1): + emptyBtn = gr.Button("Clear History") + max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True) + top_p = gr.Slider(0, 1, value=0.4, step=0.01, label="Top P", interactive=True) + temperature = gr.Slider(0, 1, value=0.8, step=0.01, label="Temperature", interactive=True) + + history = gr.State([]) + + submitBtn.click(predict, [user_input, image_path, chatbot, max_length, top_p, temperature, history], [chatbot, history], + show_progress=True) + + image_path.upload(predict_new_image, [image_path, chatbot, max_length, top_p, temperature], [chatbot, history], + show_progress=True) + image_path.clear(reset_state, outputs=[image_path, chatbot, history], show_progress=True) + + submitBtn.click(reset_user_input, [], [user_input]) + + emptyBtn.click(reset_state, outputs=[image_path, chatbot, history], show_progress=True) + +demo.queue().launch(share=False, inbrowser=True)