MMIRAGE (Modular Multimodal Intelligent Reformatting and Augmentation Generation Engine) is a framework for dataset reformatting and augmentation with language models and vision-language models.
This artifact provides the code and configuration required to reproduce the MedTrinity demo formatting pipeline used in the submitted workshop paper.
This section explains how to download the MedTrinity demo dataset and run the MMIRAGE formatting pipeline.
The reproduction uses the MedTrinity demo subset, loaded with the Hugging Face config 25M_demo. The saved local dataset contains the columns image, id, and caption. The image is stored directly inside the Hugging Face dataset, so no separate image_base_path is required.
Clone or unpack the anonymous artifact repository provided with the submission.
MMIRAGE can be run either in a Python virtual environment or inside Docker.
Clone the repository:
git clone <artifact-repo-url> MMIRAGE
cd MMIRAGEOr unpack the anonymous artifact archive:
unzip MMIRAGE-artifact.zip
cd MMIRAGECreate and activate a virtual environment:
python -m venv .venv
source .venv/bin/activateInstall MMIRAGE:
pip install --upgrade pip
pip install -e .Install the GPU extra when using the SGLang-backed llm processor for local GPU inference:
pip install -e ".[gpu]"If using GPU-backed models through SGLang, ensure that the CUDA and GPU dependencies required by MMIRAGE and SGLang are available in the environment.
The repository includes Docker and Docker Compose configurations for both GPU and CPU environments.
The docker-compose.yml defines two services:
mmirage(GPU)mmirage-cpu
The GPU container requires:
- NVIDIA GPU drivers
- NVIDIA Container Toolkit /
nvidia-container-runtime - Docker with GPU support enabled
Build and run:
docker compose build mmirage
docker compose run --rm -it mmirageThe CPU image installs MMIRAGE without GPU dependencies and can be used for workflows that do not instantiate the local SGLang-backed llm processor.
Build and run:
docker compose build mmirage-cpu
docker compose run --rm -it mmirage-cpuOpen the MedTrinity dataset page and accept the dataset access conditions:
https://huggingface.co/datasets/UCSC-VLAA/MedTrinity-25M
Then log in locally:
hf auth loginChoose a local data directory:
export DATA_ROOT=/path/to/data
export MEDTRINITY_DEMO=${DATA_ROOT}/medtrinity_demoDownload and save the demo dataset locally:
python scripts/download_medtrinity_demo.py \
--output-dir ${MEDTRINITY_DEMO}Expected output:
${MEDTRINITY_DEMO}/
dataset_info.json
state.json
data-*.arrow
medtrinity_demo_download_metadata.json
Run:
python scripts/verify_medtrinity_demo.py \
--dataset-path ${MEDTRINITY_DEMO}The verification script checks that:
- the dataset can be loaded with
datasets.load_from_disk - the required columns
image,id, andcaptionexist - images decode correctly as PIL images
- captions and ids are non-empty strings
Set the output and cache directories:
export SCRATCH=/path/to/scratch
export HF_HOME=/path/to/hfRun MMIRAGE with the local reviewer configuration:
mmirage run --config configs/medtrinity_demo_local.yaml --statsOutputs are written to:
${SCRATCH}/medtrinity_demo_conversations_formatted_local
Output schema:
{
"id": "...",
"conversations": [
{
"role": "user",
"content": "Describe this medical image."
},
{
"role": "assistant",
"content": "..."
}
],
"modalities": [
"<image stored in the dataset>"
]
}Inspect run statistics:
mmirage stats --config configs/medtrinity_demo_local.yamlSet:
export DATA_ROOT=/path/to/data
export MEDTRINITY_DEMO=${DATA_ROOT}/medtrinity_demo
export SCRATCH=/path/to/scratch
export HF_HOME=/path/to/hf
export EDF_ENV=/path/to/edfRun:
mmirage run --config configs/medtrinity_demo_16nodes.yaml --statsThe SLURM configuration uses 16 logical shards:
num_shards: 16
shard_id: "$SLURM_ARRAY_TASK_ID"Each shard requests one node:
nodes: 1
gpus: 4When all array tasks run concurrently, the pipeline uses 16 nodes total.
Inspect statistics:
mmirage stats --config configs/medtrinity_demo_16nodes.yamlconfigs/medtrinity_demo_local.yaml
configs/medtrinity_demo_16nodes.yaml
scripts/download_medtrinity_demo.py
scripts/verify_medtrinity_demo.py
MMIRAGE supports:
- text and image processing with LLMs and VLMs
- YAML-configurable pipelines
- Jinja2 prompt templating
- JMESPath-based dataset extraction
- distributed sharded execution
- local and SLURM execution modes
- configurable structured outputs
- modular processors, loaders, and writers
Run a pipeline:
mmirage run --config configs/config_mock.yamlCheck shard status:
mmirage check --config configs/config_mock.yamlRetry failed shards:
mmirage check --config configs/config_mock.yaml --retryMerge shard outputs:
mmirage merge --config configs/config_mock.yamlInspect statistics:
mmirage stats --config configs/config_mock.yamlThe following simplified configuration illustrates the main MMIRAGE sections.
processors:
- type: llm
server_args:
model_path: Qwen/Qwen2-VL-7B-Instruct
tp_size: 4
trust_remote_code: true
chat_template: qwen2-vl
default_sampling_params:
temperature: 0.1
top_p: 0.95
max_new_tokens: 768
loading_params:
state_dir: /path/to/state/dir
datasets:
- path: /path/to/image/dataset
type: loadable
output_dir: /path/to/output/shards
num_shards: 4
shard_id: "$SLURM_ARRAY_TASK_ID"
batch_size: 32
processing_params:
inputs:
- name: medical_image
key: image
type: image
- name: original_caption
key: caption
type: text
outputs:
- name: enhanced_caption
type: llm
output_type: plain
prompt: |
Describe the medical image in detail.
Original caption for context: {{ original_caption }}
remove_columns: false
output_schema:
image: "{{ medical_image }}"
caption: "{{ enhanced_caption }}"
original_caption: "{{ original_caption }}"
execution_params:
mode: local
retry: false
merge: falseMain configuration sections:
processors: model and inference configurationloading_params: dataset loading and shardingprocessing_params: prompts, extracted variables, and output schemaexecution_params: local or SLURM execution behavior
Enable runtime and throughput statistics with:
mmirage run --config configs/config_mock.yaml --statsInspect results:
mmirage stats --config configs/config_mock.yamlCommon metrics include:
runtime_secondsrows_processedthroughput_rows_per_secgpu_util_meantokens_per_sec_per_gpugpu_days_per_billion_tokens
Token metrics are null when no LLM processor is active. GPU metrics are null when nvidia-smi is unavailable or --stats was not enabled.
mmirage/
├── config/ # Configuration loading and validation
├── core/
│ ├── loader/ # Dataset loaders
│ ├── process/ # Processors and variable system
│ │ └── processors/
│ │ └── llm/ # LLM/VLM processor
│ └── writer/ # Output rendering
├── shard_process.py # Main shard processing script
└── merge_shards.py # Shard merging utility
- Jinja2: https://jinja.palletsprojects.com/en/stable/
- JMESPath: https://jmespath.org/
- SGLang: https://github.com/sgl-project/sglang
- DataTrove benchmark: https://github.com/huggingface/datatrove/tree/main/examples/inference/benchmark