Skip to content

IDEA-FinAI/ChartMoE

Repository files navigation

ChartMoE: Mixture of Expert Connector for Advanced Chart Understanding

arXiv Project Page Hugging Face Dataset

ChartMoE is a multimodal large language model with Mixture-of-Expert connector for advanced chart 1)understanding, 2)replot, 3)editing, 4)highlighting and 5)transformation.

Training Pipeline of ChartMoE

Overview

Installation

Step 1. Create a conda environment and activate it.

conda create -n chartmoe_env python=3.9
conda activate chartmoe_env

Step 2. Install PyTorch (We use PyTorch 2.1.0 / CUDA 12.1)

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu121

Step 3. Install require packages

pip install -r requirements.txt

Step 4. Install editable ChartMoE packages

pip install -e .

Quick Start

Customize the weight path of ChartMoE: Set your own ChartMoE_HF_PATH.

Code Demo:

from chartmoe import ChartMoE_Robot
import torch

robot = ChartMoE_Robot()
image_path = "examples/bar2.png"
question = "Redraw the chart with python matplotlib, giving the code to highlight the column corresponding to the year in which the student got the highest score (painting it red). Please keep the same colors and legend as the input chart."

history = ""
with torch.cuda.amp.autocast():
    response, history = robot.chat(image_path, question, history=history)

print(response)

Evaluation

ChartQA

Customize the path of ChartQA: Set your own ChartQA_ROOT(including test_human.json and test_augmented.json) and ChartQA_TEST_IMG_ROOT(including the test images).

w/ PoT:

CUDA_VISIBLE_DEVICES=0 python chartmoe/eval_ChartQA.py --save_path ./results/chartqa_results_pot --pot

w/o PoT:

CUDA_VISIBLE_DEVICES=0 python chartmoe/eval_ChartQA.py --save_path ./results/chartqa_results

MME

Run chartmoe/eval_MME.ipynb for MME scores.

WebUI Demo

CUDA_VISIBLE_DEVICES=0 python gradio_demo.py 

Acknowledgement

Thanks to InternLM-XComposer2 and CuMo for their releases of model weights and source codes!

Citation

If you find our idea or code inspiring, please cite our paper:

@article{ChartMoE,
    title={ChartMoE: Mixture of Expert Connector for Advanced Chart Understanding},
    author={Zhengzhuo Xu and Bowen Qu and Yiyan Qi and Sinan Du and Chengjin Xu and Chun Yuan and Jian Guo},
    journal={ArXiv},
    year={2024},
    volume={abs/2409.03277},
}

This code is partially based on ChartBench, if you use our code, please also cite:

@article{ChartBench,
    title={ChartBench: A Benchmark for Complex Visual Reasoning in Charts},
    author={Zhengzhuo Xu and Sinan Du and Yiyan Qi and Chengjin Xu and Chun Yuan and Jian Guo},
    journal={ArXiv},
    year={2023},
    volume={abs/2312.15915},
}