import spaces import argparse import hashlib import json import os import time from threading import Thread import logging import gradio as gr import torch from tinychart.model.builder import load_pretrained_model from tinychart.mm_utils import ( KeywordsStoppingCriteria, load_image_from_base64, process_images, tokenizer_image_token, get_model_name_from_path, ) from PIL import Image from io import BytesIO import base64 import torch from transformers import StoppingCriteria from tinychart.constants import ( DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX, ) from tinychart.conversation import SeparatorStyle, conv_templates, default_conversation from tinychart.eval.eval_metric import parse_model_output, evaluate_cmds from transformers import TextIteratorStreamer from pathlib import Path DEFAULT_MODEL_PATH = "mPLUG/TinyChart-3B-768" DEFAULT_MODEL_NAME = "TinyChart-3B-768" block_css = """ #buttons button { min-width: min(120px,100%); } """ title_markdown = """ # TinyChart: Efficient Chart Understanding with Visual Token Merging and Program-of-Thoughts Learning 🔗 [[Code](https://github.com/X-PLUG/mPLUG-DocOwl/tree/main/TinyChart)] | 📚 [[Paper](https://arxiv.org/abs/2404.16635)] **Note:** 1. Currently, this demo only supports English chart understanding and may not work well with other languages. 2. To use Program-of-Thoughts answer, please append "Answer with detailed steps." to your question. """ tos_markdown = """ ### Terms of use By using this service, users are required to agree to the following terms: The service is a research preview intended for non-commercial use only. It only provides limited safety measures and may generate offensive content. It must not be used for any illegal, harmful, violent, racist, or sexual purposes. For an optimal experience, please use desktop computers for this demo, as mobile devices may compromise its quality. """ def regenerate(state, image_process_mode): state.messages[-1][-1] = None prev_human_msg = state.messages[-2] if type(prev_human_msg[1]) in (tuple, list): prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode) state.skip_next = False return (state, state.to_gradio_chatbot(), "", None) def clear_history(): state = default_conversation.copy() return (state, state.to_gradio_chatbot(), "", None) def add_text(state, text, image, image_process_mode): if len(text) <= 0 and image is None: state.skip_next = True return (state, state.to_gradio_chatbot(), "", None) text = text[:1536] # Hard cut-off if image is not None: text = text[:1200] # Hard cut-off for images if "" not in text: # text = '' + text # text = text + "\n" text = "\n"+text text = (text, image, image_process_mode) if len(state.get_images(return_pil=True)) > 0: state = default_conversation.copy() state.append_message(state.roles[0], text) state.append_message(state.roles[1], None) state.skip_next = False return (state, state.to_gradio_chatbot(), "", None) def load_demo(): state = default_conversation.copy() return state def is_float(value): try: float(value) return True except ValueError: return False @torch.inference_mode() @spaces.GPU def get_response(params): model.to("cuda") prompt = params["prompt"] ori_prompt = prompt images = params.get("images", None) num_image_tokens = 0 if images is not None and len(images) > 0: if len(images) > 0: if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN): raise ValueError( "Number of images does not match number of tokens in prompt" ) images = [load_image_from_base64(image) for image in images] images = process_images(images, image_processor, model.config) if type(images) is list: images = [ image.to(model.device, dtype=torch.float16) for image in images ] else: images = images.to(model.device, dtype=torch.float16) replace_token = DEFAULT_IMAGE_TOKEN if getattr(model.config, "mm_use_im_start_end", False): replace_token = ( DEFAULT_IM_START_TOKEN + replace_token + DEFAULT_IM_END_TOKEN ) prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) if hasattr(model.get_vision_tower().config, "tome_r"): num_image_tokens = ( prompt.count(replace_token) * model.get_vision_tower().num_patches - 26 * model.get_vision_tower().config.tome_r ) else: num_image_tokens = ( prompt.count(replace_token) * model.get_vision_tower().num_patches ) else: images = None image_args = {"images": images} else: images = None image_args = {} temperature = float(params.get("temperature", 1.0)) top_p = float(params.get("top_p", 1.0)) max_context_length = getattr(model.config, "max_position_embeddings", 2048) max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024) stop_str = params.get("stop", None) do_sample = True if temperature > 0.001 else False logger.info(prompt) input_ids = ( tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt") .unsqueeze(0) .to(model.device) ) keywords = [stop_str] stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) streamer = TextIteratorStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=1500 ) max_new_tokens = min( max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens ) if max_new_tokens < 1: yield json.dumps( { "text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0, } ).encode() + b"\0" return # local inference # BUG: If stopping_criteria is set, an error occur: # RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 0 generate_kwargs = dict( inputs=input_ids, do_sample=do_sample, temperature=temperature, top_p=top_p, max_new_tokens=max_new_tokens, streamer=streamer, # stopping_criteria=[stopping_criteria], use_cache=True, **image_args, ) thread = Thread(target=model.generate, kwargs=generate_kwargs) thread.start() logger.debug(ori_prompt) logger.debug(generate_kwargs) generated_text = ori_prompt for new_text in streamer: generated_text += new_text if generated_text.endswith(stop_str): generated_text = generated_text[: -len(stop_str)] yield json.dumps({"text": generated_text, "error_code": 0}).encode() if '' in generated_text and '' in generated_text and '' in generated_text and '' in generated_text: program = generated_text program = '#' + program.split('ASSISTANT: #')[-1] print(program) try: execuate_result = evaluate_cmds(parse_model_output(program)) if is_float(execuate_result): execuate_result = round(float(execuate_result), 4) generated_text += f'\n\nExecute result: {execuate_result}' yield json.dumps({"text": generated_text, "error_code": 0}).encode() + b"\0" except: # execuate_result = 'Failed.' generated_text += f'\n\nIt seems the execution of the above code encounters bugs. I\'m trying to answer this question directly...' ori_generated_text = generated_text + '\nDirect Answer: ' direct_prompt = ori_prompt.replace(' Answer with detailed steps.', '') direct_input_ids = ( tokenizer_image_token(direct_prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt") .unsqueeze(0) .to(model.device) ) generate_kwargs = dict( inputs=direct_input_ids, do_sample=do_sample, temperature=temperature, top_p=top_p, max_new_tokens=max_new_tokens, streamer=streamer, use_cache=True, **image_args, ) thread = Thread(target=model.generate, kwargs=generate_kwargs) thread.start() generated_text = ori_generated_text for new_text in streamer: generated_text += new_text if generated_text.endswith(stop_str): generated_text = generated_text[: -len(stop_str)] yield json.dumps({"text": generated_text, "error_code": 0}).encode() def http_bot(state, temperature, top_p, max_new_tokens): if state.skip_next: # This generate call is skipped due to invalid inputs yield (state, state.to_gradio_chatbot()) return if len(state.messages) == state.offset + 2: # First round of conversation template_name = 'phi' new_state = conv_templates[template_name].copy() new_state.append_message(new_state.roles[0], state.messages[-2][1]) new_state.append_message(new_state.roles[1], None) state = new_state # Construct prompt prompt = state.get_prompt() all_images = state.get_images(return_pil=True) all_image_hash = [hashlib.md5(image.tobytes()).hexdigest() for image in all_images] # Make requests # pload = {"model": model_name, "prompt": prompt, "temperature": float(temperature), "top_p": float(top_p), # "max_new_tokens": min(int(max_new_tokens), 1536), "stop": ( # state.sep # if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] # else state.sep2 # ), "images": state.get_images()} pload = { "model": model_name, "prompt": prompt, "temperature": float(temperature), "top_p": float(top_p), "max_new_tokens": min(int(max_new_tokens), 1536), "stop": ( state.sep if state.sep_style in [SeparatorStyle.SINGLE, SeparatorStyle.MPT] else state.sep2 ), "images": state.get_images()} state.messages[-1][-1] = "▌" yield (state, state.to_gradio_chatbot()) # for stream output = get_response(pload) for chunk in output: if chunk: data = json.loads(chunk.decode().replace('\x00','')) if data["error_code"] == 0: output = data["text"][len(prompt) :].strip() state.messages[-1][-1] = output + "▌" yield (state, state.to_gradio_chatbot()) else: output = data["text"] + f" (error_code: {data['error_code']})" state.messages[-1][-1] = output yield (state, state.to_gradio_chatbot()) return time.sleep(0.03) state.messages[-1][-1] = state.messages[-1][-1][:-1] yield (state, state.to_gradio_chatbot()) def build_demo(): textbox = gr.Textbox( show_label=False, placeholder="Enter text and press ENTER", container=False ) with gr.Blocks(title="TinyChart", theme=gr.themes.Default(), css=block_css) as demo: state = gr.State() gr.Markdown(title_markdown) with gr.Row(): with gr.Column(scale=5): with gr.Row(elem_id="Model ID"): gr.Dropdown( choices=[DEFAULT_MODEL_NAME], value=DEFAULT_MODEL_NAME, interactive=True, label="Model ID", container=False, ) imagebox = gr.Image(type="pil") image_process_mode = gr.Radio( ["Crop", "Resize", "Pad", "Default"], value="Default", label="Preprocess for non-square image", visible=False, ) cur_dir = Path(__file__).parent gr.Examples( examples=[ [ f"{cur_dir}/images/market.png", "What is the highest number of companies in the domestic market? Answer with detailed steps.", ], [ f"{cur_dir}/images/college.png", "What is the difference between Asians and Whites degree distribution? Answer with detailed steps." ], [ f"{cur_dir}/images/immigrants.png", "How many immigrants are there in 1931?", ], [ f"{cur_dir}/images/sails.png", "By how much percentage wholesale is less than retail? Answer with detailed steps." ], [ f"{cur_dir}/images/diseases.png", "Is the median value of all the bars greater than 30? Answer with detailed steps.", ], [ f"{cur_dir}/images/economy.png", "Which team has higher economy in 28 min?" ], [ f"{cur_dir}/images/workers.png", "Generate underlying data table for the chart." ], [ f"{cur_dir}/images/sports.png", "Create a brief summarization or extract key insights based on the chart image." ], [ f"{cur_dir}/images/albums.png", "Redraw the chart with Python code." ] ], inputs=[imagebox, textbox], ) with gr.Accordion("Parameters", open=False) as _: temperature = gr.Slider( minimum=0.0, maximum=1.0, value=0.1, step=0.1, interactive=True, label="Temperature", ) top_p = gr.Slider( minimum=0.0, maximum=1.0, value=0.7, step=0.1, interactive=True, label="Top P", ) max_output_tokens = gr.Slider( minimum=0, maximum=1024, value=1024, step=64, interactive=True, label="Max output tokens", ) with gr.Column(scale=8): chatbot = gr.Chatbot(elem_id="chatbot", label="Chatbot", height=550) with gr.Row(): with gr.Column(scale=8): textbox.render() with gr.Column(scale=1, min_width=50): submit_btn = gr.Button(value="Send", variant="primary") with gr.Row(elem_id="buttons") as _: regenerate_btn = gr.Button(value="🔄 Regenerate", interactive=True) clear_btn = gr.Button(value="🗑️ Clear", interactive=True) gr.Markdown(tos_markdown) regenerate_btn.click( regenerate, [state, image_process_mode], [state, chatbot, textbox, imagebox], queue=False, ).then( http_bot, [state, temperature, top_p, max_output_tokens], [state, chatbot] ) clear_btn.click( clear_history, None, [state, chatbot, textbox, imagebox], queue=False ) textbox.submit( add_text, [state, textbox, imagebox, image_process_mode], [state, chatbot, textbox, imagebox], queue=False, ).then( http_bot, [state, temperature, top_p, max_output_tokens], [state, chatbot] ) submit_btn.click( add_text, [state, textbox, imagebox, image_process_mode], [state, chatbot, textbox, imagebox], queue=False, ).then( http_bot, [state, temperature, top_p, max_output_tokens], [state, chatbot] ) demo.load(load_demo, None, [state], queue=False) return demo def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("--host", type=str, default=None) parser.add_argument("--port", type=int, default=None) parser.add_argument("--share", default=None) parser.add_argument("--model-path", type=str, default=DEFAULT_MODEL_PATH) parser.add_argument("--model-name", type=str, default=DEFAULT_MODEL_NAME) parser.add_argument("--load-8bit", action="store_true") parser.add_argument("--load-4bit", action="store_true") args = parser.parse_args() return args if __name__ == "__main__": logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", ) logger = logging.getLogger(__name__) logger.info(gr.__version__) args = parse_args() model_name = args.model_name tokenizer, model, image_processor, context_len = load_pretrained_model( model_path=args.model_path, model_base=None, model_name=args.model_name, device="cpu", load_4bit=args.load_4bit, load_8bit=args.load_8bit, torch_dtype=torch.float16, ) demo = build_demo() demo.queue() demo.launch(server_name=args.host, server_port=args.port, share=args.share)