Skip to content

README.md

Latest commit

 

History

History
282 lines (230 loc) · 9.05 KB

README.md

File metadata and controls

282 lines (230 loc) · 9.05 KB

MMIRAGE

MMIRAGE, which stands for Modular Multimodal Intelligent Reformatting and Augmentation Generation Engine, is an advanced platform designed to streamline the processing of datasets using generative models, including vision-language models (VLMs). It is engineered to handle large-scale data reformatting and augmentation tasks with efficiency and precision. By leveraging state-of-the-art generative models, MMIRAGE enables users to perform complex dataset transformations, ensuring compatibility across various formats and schemas. Its multi-node support and parallel processing capabilities make it an ideal choice for scenarios demanding substantial computational power, such as distributed training and inference workflows. MMIRAGE not only simplifies the integration of powerful language models but also provides a customizable framework for diverse use cases, from reformatting conversational datasets to generating Q/A pairs from plain text.

How to install

To install the library, you can clone it from GitHub and then use pip to install it directly. It is recommended to have already installed torch and sglang to take advantage of GPU acceleration.

git clone git@github.com:EPFLiGHT/MMIRAGE.git
pip install -e ./MMIRAGE

For testing and scripts that make use of the library, it is advised to create a .env file:

./scripts/generate_env.sh

Key features

  • Multimodal Support: Process both text and images with vision-language models
  • Easily configurable with a YAML file which configures the following parameters:
    • The prompt to the LLM (using Jinja2 templating)
    • Variables with the name and their JMESPath key to a JSON
    • Image inputs for multimodal processing
  • Parallelizable with multi-node support
    • The training pipeline uses distributed inference with sharding
  • Support a variety of LLMs, VLMs (Vision-Language Models), and image generation models
  • Support any dataset schemas (configurable with the YAML format)
  • The ability to either output a JSON (or any other structured format) or plain text
  • Modular architecture with pluggable processors, loaders, and writers

Example usage

Running (single command)

Run the pipeline via the CLI. Retry behavior is driven by your YAML config:

  • execution_params.retry: true → automatically retries failed shards until completion or max_retries
  • execution_params.retry: false → submits/runs once; you can later trigger retries via check
mmirage run --config configs/config_mock.yaml

To check status only:

mmirage check --config configs/config_mock.yaml

To check status and submit retries for failed shards:

mmirage check --config configs/config_mock.yaml --retry

Text-only: Reformatting dataset

Suppose you have a dataset with samples of the following format

{ 
    "conversations" : [{"role": "user", "content": "Describe the image"}, {"role": "assistant", "content": "This is a badly formmatted answer"}],
    "modalities" : ["<the images>"]
}

The dataset contains assistant answers that are badly formatted. The goal would be to use a LLM to format our answer in Markdown. With MMIRAGE, it would be as simple as defining a YAML configuration file:

processors:
  - type: llm
    server_args:
      model_path: Qwen/Qwen3-8B
      tp_size: 4
      trust_remote_code: true
    default_sampling_params:
      temperature: 0.1
      top_p: 1.0
      max_new_tokens: 384

loading_params:
  state_dir: /path/to/state/dir
  datasets:
    - path: /path/to/dataset
      type: loadable
      output_dir: /path/to/output/shards
  num_shards: 4
  shard_id: "$SLURM_ARRAY_TASK_ID"
  batch_size: 64

processing_params:
  inputs:
    - name: assistant_answer
      key: conversations[1].content
    - name: user_prompt
      key: conversations[0].content
    - name: modalities
      key: modalities

  outputs:
    - name: formatted_answer
      type: llm
      output_type: plain
      prompt: | 
        Reformat the answer in a markdown format without adding anything else:
        {{ assistant_answer }}
      
  remove_columns: false
  output_schema:
    conversations:
      - role: user
        content: "{{ user_prompt }}"
      - role: assistant
        content: "{{ formatted_answer }}"
    modalities: "{{ modalities }}"

execution_params:
  mode: local
  retry: false

Configuration explanation:

  • processors: List of processor configurations. Supports llm (text/VLM generation) and image_gen (text-to-image generation).
  • loading_params: Parameters for loading and sharding datasets.
    • state_dir: Optional shared directory for shard status/retry state. Defaults to ~/.cache/MMIRAGE/state_dir.
    • datasets: List of dataset configurations with path, type, and output directory.
  • processing_params:
    • inputs: Variables extracted from the input dataset using JMESPath queries.
    • outputs: Variables created by processors. Prompts use Jinja2 templating ({{ variable }}).
    • output_schema: Defines the structure of output samples.
  • execution_params:
    • mode: "local" to run shard processing in the current Python environment or "slurm" to run through SLURM by submitting an sbatch array job.
    • retry: If true, MMIRAGE automatically retries failed shards until they succeed or max_retries is reached. If false, the pipeline runs/submits once, and retries can be triggered later via the check/retry CLI commands.

Multimodal: Processing images with VLMs

MMIRAGE supports multimodal processing with vision-language models:

processors:
  - type: llm
    server_args:
      model_path: Qwen/Qwen2-VL-7B-Instruct
      tp_size: 4
      trust_remote_code: true
    chat_template: qwen2-vl  # Required for VLMs
    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  # Mark as image input
      image_base_path: /path/to/images  # Base directory for relative paths
    - 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

Image generation: Text-to-image pipeline

MMIRAGE also supports image generation with Diffusers models:

processors:
  - type: image_gen
    pipeline_args:
      model_path: stable-diffusion-v1-5/stable-diffusion-v1-5
      torch_dtype: float16
      device: auto
      enable_attention_slicing: true
    default_sampling_params:
      num_inference_steps: 20
      guidance_scale: 7.5
    output_dir: /path/to/generated/images
    file_format: png

loading_params:
  state_dir: /path/to/state/dir
  datasets:
    - path: /path/to/prompts.jsonl
      type: JSONL
      output_dir: /path/to/output/shards
  num_shards: 1
  shard_id: 0
  batch_size: 8

processing_params:
  inputs:
    - name: prompt_text
      key: text

  outputs:
    - name: generated_image
      type: image_gen
      output_mode: path          # "path" or "pil"
      filename_template: "generated_{{ __shard_id }}_{{ __sample_index }}_{{ __source_hash }}"
      width: 512
      height: 512
      prompt: |
        Create an illustration of:
        {{ prompt_text }}

  remove_columns: false
  output_schema:
    text: "{{ prompt_text }}"
    image: "{{ generated_image }}"

execution_params:
  mode: local
  retry: false

Install optional image generation dependencies before running this config:

pip install -e .[image_gen]

Key multimodal features:

  • chat_template: Specify the VLM chat template (e.g., qwen2-vl)
  • type: image: Mark input variables as images
  • image_base_path: Base directory for resolving relative image paths
  • Supports PIL Images, URLs, and file paths

Architecture

MMIRAGE uses a modular architecture:

mmirage/
├── config/           # Configuration loading and validation
├── core/
│   ├── loader/       # Dataset loaders (JSONL, HuggingFace)
│   ├── process/      # Processors (LLM, etc.) and variable system
│   │   └── processors/
│   │       └── llm/  # LLM processor with multimodal support
│   └── writer/       # Output rendering with Jinja2
├── shard_process.py  # Main processing script
└── merge_shards.py   # Shard merging utility

Useful tools

  • Jinja2 for template processing: link
  • JMESPath for JSON queries: link
  • SGLang for fast inference: link
  • Performance paper: link