diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..cce1e56 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,26 @@ +# CI for NV-Reason-CXR: lint and format checks via pre-commit +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + pre-commit: + name: pre-commit + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Install pre-commit + run: pip install pre-commit + + - name: Run pre-commit + run: pre-commit run --all-files diff --git a/.markdownlint.yaml b/.markdownlint.yaml new file mode 100644 index 0000000..998dbbd --- /dev/null +++ b/.markdownlint.yaml @@ -0,0 +1,7 @@ +# Markdownlint config for NV-Generate-CTMR +# Relaxed for existing docs (READMEs with tables, HTML, long lines). +# Re-enable rules as you clean up docs or for new files. + +# Line length: allow long lines common in docs (tables, code, links) +MD013: + line_length: 700 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..5a66a9e --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,35 @@ +# Pre-commit hooks for NV-Reason-CXR +# Install: pip install pre-commit && pre-commit install +# Run manually: pre-commit run --all-files + +repos: + # General file checks + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + - id: check-json + - id: check-merge-conflict + - id: check-added-large-files + args: [--maxkb=1000] + - id: check-case-conflict + - id: debug-statements + + # Python linting and formatting (ruff) — fixes applied locally + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.8.4 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + # Markdown linting + - repo: https://github.com/igorshubovych/markdownlint-cli + rev: v0.38.0 + hooks: + - id: markdownlint + +ci: + autoupdate_commit_msg: "chore: pre-commit autoupdate" diff --git a/README.md b/README.md index 212fc6b..e719404 100644 --- a/README.md +++ b/README.md @@ -1,19 +1,20 @@ # NV-Reason-CXR-3B ## Description + NV-Reason-CXR-3B is a specialized vision-language model designed for medical reasoning and interpretation of chest X-ray images, with detailed explanations. The model combines visual understanding with medical reasoning capabilities, enabling healthcare professionals to access comprehensive analyses and engage in follow-up discussions about radiological findings. NV-Reason-CXR-3B provides step-by-step reasoning that mirrors clinical thinking patterns, making it valuable for educational and research applications in medical imaging. This model is for research and development only. It is intended to empower developers to extend this work in their tasks and to provide practical examples of applying the methodology across medical domains. -**Table of Contents** -1. [Overview](#overview) -2. [Introduction](#introduction) -3. [Installation](#installation) -4. [Training models](#training-models) - - [SFT](#sft) - - [GRPO](#grpo) -5. [Data](#data) +## Table of Contents +1. [Overview](#overview) +2. [Introduction](#introduction) +3. [Installation](#installation) +4. [Training models](#training-models) + - [SFT](#sft) + - [GRPO](#grpo) +5. [Data](#data) ## Overview @@ -25,37 +26,35 @@ The goal of this repo is to provide examples for inference and training of the [ ## Introduction -Vision–language models (VLMs) have shown strong promise for medical image analysis, but most remain opaque, offering predictions without the transparent, stepwise reasoning clinicians rely on. We present a framework that brings chain-of-thought (CoT) reasoning to chest X-ray interpretation. +Vision–language models (VLMs) have shown strong promise for medical image analysis, but most remain opaque, offering predictions without the transparent, stepwise reasoning clinicians rely on. We present a framework that brings chain-of-thought (CoT) reasoning to chest X-ray interpretation. Our approach is designed to learn how experts reason—not just what they conclude—by aligning intermediate steps with observable image evidence and radiology workflow. Beyond accuracy, the explicit reasoning traces support clinical auditability: they reveal why a conclusion was reached, which alternatives were considered, and where uncertainty remains—enabling quality assurance, error analysis, and safer human–AI collaboration. Inspired by reasoning-first training (DeepSeek-R1 and Open-R1), our approach combines a radiologist-style supervised fine-tuning (SFT) warm start with GRPO reinforcement learning (RL) and verifiable rewards defined over a list of chest X-ray abnormalities. -We enlisted several experienced radiologists to annotate their internal reasoning while reading chest X-ray cases. To support this, we developed an internal web platform that makes thought capture as seamless as possible. The platform provides automated voice recording, transcription, error correction, and optional translation into English. We used both the collected human reasoning data and synthetic reasoning data for training with SFT, as well as abnormality list only (from MIMIC-CXR) for GRPO training. +We enlisted several experienced radiologists to annotate their internal reasoning while reading chest X-ray cases. To support this, we developed an internal web platform that makes thought capture as seamless as possible. The platform provides automated voice recording, transcription, error correction, and optional translation into English. We used both the collected human reasoning data and synthetic reasoning data for training with SFT, as well as abnormality list only (from MIMIC-CXR) for GRPO training. In an expert reader study, AI-assisted reasoning increased confidence, supported targeted error auditing, and reduced time to finalize reports—particularly for abnormal cases. On out-of-distribution (OOD) evaluation using the CheXpert test set, the model attains competitive multi-label classification while providing faithful rationales. NV-Reason-CXR-3B is designed to respond in the style of a teacher, a senior radiologist, explaining the problem and the solution and offers: -- Chain-of-thought processing - - The reasoning engine generates step-by-step diagnostic analysis - - Systematic anatomical review - - Identification of normal and abnormal findings - - Differential diagnosis consideration -- Clinical output generation - - Main findings - - Step-by-step reasoning pathway - - Differential diagnoses and their likelihood - - Recommendations for follow-up or clinical correlation - - Clarification multi-step follow-up chat - - Structured report generation - +- Chain-of-thought processing + - The reasoning engine generates step-by-step diagnostic analysis + - Systematic anatomical review + - Identification of normal and abnormal findings + - Differential diagnosis consideration +- Clinical output generation + - Main findings + - Step-by-step reasoning pathway + - Differential diagnoses and their likelihood + - Recommendations for follow-up or clinical correlation + - Clarification multi-step follow-up chat + - Structured report generation An example of the model output: -![](docs/d1.png) - -You can try the 🩻 [\[Web Demo\]](https://huggingface.co/spaces/nvidia/nv-reason-cxr) for examples of the model output, where you can also ask the follow up questions such as "provide differentials" and "write a structured report". +![Model output example](docs/d1.png) +You can try the 🩻 [\[Web Demo\]](https://huggingface.co/spaces/nvidia/nv-reason-cxr) for examples of the model output, where you can also ask the follow up questions such as "provide differentials" and "write a structured report". ## Preliminary subjective evaluation @@ -69,27 +68,26 @@ text); - Full AI reasoning: Full AI reasoning output and the structured report. Readers were instructed to behave as in routine practice. - | ![1](docs/se1.png) | ![2](docs/se2.png) | | - | - | | ![3](docs/se3.png) | ![4](docs/se4.png) | Overall, experts rated the reasoning traces as accurate, appropriately qualified, and practically useful; full reasoning notably improved trust and confidence and yielded substantial time savings—especially for abnormal studies. +### Use Case -### Use Case: Radiologists, medical students, and medical researchers would be expected to use this system for chest X-ray interpretation with detailed reasoning, educational training with AI-generated explanations, and research applications requiring explainable medical AI analyses. **Important Medical AI Considerations:** This model is designed for research and educational purposes only and should not be used for clinical diagnosis or treatment decisions. All outputs should be reviewed by qualified medical professionals. The model's reasoning capabilities are intended to support medical education and research, not replace clinical judgment. -## Model Architecture: +## Model Architecture + - **Architecture Type:** Transformer - **Network Architecture:** Vision-Language Model based on [Qwen2.5-VL-3B](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) architecture with medical reasoning capabilities This model was developed by fine-tuning Qwen2.5-VL-3B using Supervised Fine-Tuning (SFT) and Group Relative Policy Optimization (GRPO) for enhanced medical reasoning. - ## Quick start / Inference ```python @@ -98,7 +96,7 @@ from transformers import AutoModelForImageTextToText, AutoProcessor from PIL import Image -# Load the model +# Load the model model_name = "nvidia/NV-Reason-CXR-3B" model = AutoModelForImageTextToText.from_pretrained( model_name, @@ -135,7 +133,7 @@ text = processor.apply_chat_template(messages, add_generation_prompt=True) inputs = processor(text=text, images=[image], return_tensors="pt") inputs = inputs.to(model.device) -# Generate +# Generate generated_ids = model.generate(**inputs, max_new_tokens=2048) # Trim and decode @@ -155,10 +153,10 @@ print(generated_text) ``` - ## Installation -For inference only, the minimal set of required dependencies are +For inference only, the minimal set of required dependencies are + ```shell pip install torch==2.7.1 torchvision==0.22.1 transformers==4.56.1 ``` @@ -167,7 +165,7 @@ For training, we recommend to create a Python virtual environment with `uv`. To install `uv`, follow instructions [here](https://docs.astral.sh/uv/getting-started/installation/). ```shell -uv venv --seed --python 3.11 nvreasoncxr && source nvreasoncxr/bin/activate +uv venv --seed --python 3.11 nvreasoncxr && source nvreasoncxr/bin/activate ``` Then, install dependencies: @@ -176,7 +174,7 @@ Then, install dependencies: uv pip install vllm==0.10.1.1 uv pip install flash-attn==2.8.3 --no-build-isolation uv pip install accelerate bitsandbytes datasets peft wandb deepspeed einops flake8 hf_transfer huggingface-hub isort liger-kernel packaging parameterized safetensors pandas numpy scikit-learn qwen-vl-utils -uv pip install trl==0.22.2 transformers==4.56.1 +uv pip install trl==0.22.2 transformers==4.56.1 ``` @@ -188,12 +186,10 @@ Optionally, log into your WANDB account to view training progress later: wandb login ``` - ## Training models The training configuration assumes a node of 8 x A100s NVIDIA GPUs (80GB). You'll need to download the images for the examples first, and place them into the "images" folder, see the "Data" section below. - ### SFT ```shell @@ -207,11 +203,10 @@ accelerate launch --config_file accelerate/zero2.yaml \ --output_dir data/output_sft_model \ --dataset_path datalists/sft.jsonl \ --num_train_epochs 1 \ - --dataset_streaming false \ + --dataset_streaming false \ --gradient_accumulation_steps 8 ``` - ### GRPO ```shell @@ -228,24 +223,22 @@ accelerate launch --config_file accelerate/zero2.yaml \ --gradient_accumulation_steps 8 ``` +## Data - -## Data The model was trained on both internally collected human reasoning data and synthetic data. The small datasets provided below are examples only, intended to demonstrate the training code. The full training dataset is currently not provided. ### Data for SFT training example -Download the x-ray images of the MIMIC-CXR-JPG dataset from [here](https://physionet.org/content/mimic-cxr-jpg/2.1.0/). You'll need to comply with the data Terms and Conditions. The training example uses only a small subset of 256 cases, so you could download only the images listed [here](datalists/sft.jsonl). In these examples the radiology thinking process was synthetically generated with LLM by rewriting the x-ray report text. This small subset is intended only as an example. Extract the image files (ignoring any subfolders) into the "images/mimic-cxr-jpg/images_512" folder. +Download the x-ray images of the MIMIC-CXR-JPG dataset from [here](https://physionet.org/content/mimic-cxr-jpg/2.1.0/). You'll need to comply with the data Terms and Conditions. The training example uses only a small subset of 256 cases, so you could download only the images listed [here](datalists/sft.jsonl). In these examples the radiology thinking process was synthetically generated with LLM by rewriting the x-ray report text. This small subset is intended only as an example. Extract the image files (ignoring any subfolders) into the "images/mimic-cxr-jpg/images_512" folder. ### Data for GRPO training example -Download the x-ray images of the test set of CheXpert dataset from [here](https://stanfordaimi.azurewebsites.net/datasets/23c56a0d-15de-405b-87c8-99c30138950c). You'll need to comply with the data Terms and Conditions. Extract the train subset (CheXpert/test) into "images/CheXpert/test" directory of this repo. +Download the x-ray images of the test set of CheXpert dataset from [here](https://stanfordaimi.azurewebsites.net/datasets/23c56a0d-15de-405b-87c8-99c30138950c). You'll need to comply with the data Terms and Conditions. Extract the train subset (CheXpert/test) into "images/CheXpert/test" directory of this repo. -We provide a data manifest [file](datalists/grpo.jsonl) formatted for GRPO training. It lists image names and solutions for each case, where the "solution" is a list of abnormalities present in each image. The task of GRPO training is to learn the thinking process based solely on the provided list of abnormalities. +We provide a data manifest [file](datalists/grpo.jsonl) formatted for GRPO training. It lists image names and solutions for each case, where the "solution" is a list of abnormalities present in each image. The task of GRPO training is to learn the thinking process based solely on the provided list of abnormalities. This CheXpert test set was used for testing during the model development. But here, instead we use it for the training example, since the data subset is very small, and you should observe accuracy improvement quickly (check WANDB graphs of accuracy). - ## Acknowledgements This project uses a number of Huggingface libraries, including TRL, Transformers and Accelerate, as well as implementation ideas from a great [open-r1](https://github.com/huggingface/open-r1) project: "Open R1: A fully open reproduction of DeepSeek-R1", Hugging Face, Jan 2025. @@ -257,18 +250,17 @@ This project uses a number of Huggingface libraries, including TRL, Transformers - **Base Model**: [Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct) - **Datasets**: [MIMIC-CXR](https://physionet.org/content/mimic-cxr-jpg/2.1.0/) • [CheXpert](https://stanfordaimi.azurewebsites.net/datasets/8cbd9ed4-2eb9-4565-affc-111cf4f7ebe2) - ## License NV-Reason-CXR-3B model weights are released under the [NVIDIA OneWay Noncommercial License Agreement](https://huggingface.co/nvidia/NV-Reason-CXR-3B/blob/main/LICENSE). - ## Citation If you find our work helpful, please consider citing the [paper](https://arxiv.org/abs/2510.23968): + ```bibtex @misc{myronenko2025reasoning, - title={Reasoning Visual Language Model for Chest X-Ray Analysis}, + title={Reasoning Visual Language Model for Chest X-Ray Analysis}, author={Andriy Myronenko and Dong Yang and Baris Turkbey and Mariam Aboian and Sena Azamat and Esra Akcicek and Hongxu Yin and Pavlo Molchanov and Marc Edgar and Yufan He and Pengfei Guo and Yucheng Tang and Daguang Xu}, year={2025}, eprint={2510.23968}, @@ -278,5 +270,3 @@ If you find our work helpful, please consider citing the [paper](https://arxiv.o url={https://arxiv.org/abs/2510.23968} } ``` - - diff --git a/configs/grpo_config.yaml b/configs/grpo_config.yaml index 2d7d023..fcbba86 100644 --- a/configs/grpo_config.yaml +++ b/configs/grpo_config.yaml @@ -1,16 +1,16 @@ # ENV vars env: - WANDB_PROJECT : nv-reason-cxr + WANDB_PROJECT : nv-reason-cxr TOKENIZERS_PARALLELISM: false PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True,max_split_size_mb:2048" - VLLM_DISABLE_COMPILE_CACHE: 1 + VLLM_DISABLE_COMPILE_CACHE: 1 # Data training arguments dataset_name: nv-reason-cxr-data dataset_path: grpo.jsonl dataset_root: datalists/ image_dir: images/ -dataset_streaming: false +dataset_streaming: false ########################################################################## @@ -26,23 +26,23 @@ attn_implementation: flash_attention_2 ############################################ # GRPO trainer config -reward_funcs: ["accuracy_reward_hard", "soft_overshort_punishment"] -# reward_weights: [1.0, 1.0] +reward_funcs: ["accuracy_reward_hard", "soft_overshort_punishment"] +# reward_weights: [1.0, 1.0] -loss_type: grpo +loss_type: grpo scale_rewards: group -temperature: 1 +temperature: 1 mask_truncated_completions: true -use_liger_loss: false +use_liger_loss: false use_vllm: true -vllm_mode: colocate -vllm_gpu_memory_utilization: 0.3 +vllm_mode: colocate +vllm_gpu_memory_utilization: 0.3 -max_prompt_length: null -max_completion_length: 1800 +max_prompt_length: null +max_completion_length: 1800 -num_generations: 16 +num_generations: 16 per_device_train_batch_size: 16 gradient_accumulation_steps: 8 per_device_eval_batch_size: 1 @@ -65,14 +65,14 @@ lr_scheduler_kwargs: warmup_steps: 5 max_grad_norm: 0.2 -eval_steps: 100 -do_eval: false +eval_steps: 100 +do_eval: false eval_strategy: "no" -save_steps: 100 -save_total_limit: 3 -metric_for_best_model: eval_loss -greater_is_better: false +save_steps: 100 +save_total_limit: 3 +metric_for_best_model: eval_loss +greater_is_better: false bf16: true @@ -86,10 +86,9 @@ seed: 123 report_to: wandb log_level: info -logging_steps: 1 +logging_steps: 1 logging_strategy: steps -logging_first_step: true +logging_first_step: true remove_unused_columns: false #VLM specific dataloader_num_workers: 4 - diff --git a/configs/sft_config.yaml b/configs/sft_config.yaml index 9012f6a..159b8b1 100644 --- a/configs/sft_config.yaml +++ b/configs/sft_config.yaml @@ -1,7 +1,7 @@ # ENV vars env: - WANDB_PROJECT : nv-reason-cxr + WANDB_PROJECT : nv-reason-cxr # Data training arguments dataset_name: nv-reason-cxr-data @@ -36,16 +36,16 @@ gradient_checkpointing_kwargs: use_reentrant: false optim: adamw_torch_fused -learning_rate: 2.0e-05 +learning_rate: 2.0e-05 lr_scheduler_type: cosine warmup_ratio: 0.03 max_grad_norm: 0.3 -do_eval: false +do_eval: false eval_strategy: "no" -save_steps: 100 -save_total_limit: 3 +save_steps: 100 +save_total_limit: 3 bf16: true tf32: true @@ -59,12 +59,12 @@ seed: 123 # run_name: run_name #defaults to output folder name report_to: wandb log_level: info -logging_steps: 20 +logging_steps: 20 logging_strategy: steps #VLM specific -remove_unused_columns: false +remove_unused_columns: false dataset_text_field: "" -dataset_kwargs: {"skip_prepare_dataset": True} -dataset_num_proc: 4 +dataset_kwargs: {"skip_prepare_dataset": True} +dataset_num_proc: 4 max_length: null diff --git a/images/README.md b/images/README.md index 42971ea..87cae08 100644 --- a/images/README.md +++ b/images/README.md @@ -1 +1,3 @@ -This folder is a placeholder to download x-ray images into, when following the training examples. \ No newline at end of file +# Images + +This folder is a placeholder to download x-ray images into, when following the training examples. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..a0254ab --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,13 @@ +# Minimal config for tooling (pre-commit, ruff). NV-Generate-CTMR has no installable package. + +[tool.ruff] +target-version = "py311" +line-length = 150 +exclude = [".git", "__pycache__", "data", "figures", "assets", "*.ipynb"] + +[tool.ruff.lint] +select = ["E", "F", "I", "N", "W", "UP"] +ignore = ["E501"] # line length handled by formatter + +[tool.ruff.format] +quote-style = "double" diff --git a/train/vlm_grpo_train.py b/train/vlm_grpo_train.py index b96245b..0ff437b 100644 --- a/train/vlm_grpo_train.py +++ b/train/vlm_grpo_train.py @@ -1,95 +1,89 @@ import logging import os import sys -from typing import Dict, Any, List -from dataclasses import dataclass, field, asdict +from dataclasses import dataclass, field +from typing import Any -import torch import datasets -from datasets import Value, disable_caching - +import torch import transformers -from transformers import set_seed, AutoModelForImageTextToText +from datasets import disable_caching +from qwen_vl_utils import process_vision_info +from transformers import AutoModelForImageTextToText, set_seed from transformers.trainer_utils import get_last_checkpoint +from trl import GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser +from vlm_rewards import accuracy_reward, accuracy_reward_hard, format_reward, get_soft_overshort_punishment, tag_count_reward from accelerate import PartialState -from trl import ModelConfig, ScriptArguments, TrlParser, GRPOTrainer, GRPOConfig -from vlm_rewards import ( - accuracy_reward, - format_reward, - tag_count_reward, - accuracy_reward_hard, - get_soft_overshort_punishment -) - -from qwen_vl_utils import process_vision_info -disable_caching() +disable_caching() logger = logging.getLogger(__name__) + @dataclass class VLMScriptArguments(ScriptArguments): - ''' + """ Additional command line arguments for the GRPO training script. For a full list of arguments, see the cofings/grpo_config.yaml file. - ''' + """ - reward_funcs: list[str] = field(default_factory=lambda: ["accuracy_reward_hard", "soft_overshort_punishment"], metadata={"help": "List of reward functions."}) + reward_funcs: list[str] = field( + default_factory=lambda: ["accuracy_reward_hard", "soft_overshort_punishment"], metadata={"help": "List of reward functions."} + ) dataset_path: str = field(default="grpo.jsonl", metadata={"help": "Path to the dataset json file."}) dataset_root: str = field(default="datalists/", metadata={"help": "Path to the datalists root directory."}) - image_dir: str = field(default="images/", metadata={"help": "Path to the image directory."}) - + image_dir: str = field(default="images/", metadata={"help": "Path to the image directory."}) + max_image_height: int = field(default=476, metadata={"help": "Max height e.g 512."}) min_image_height: int = field(default=128, metadata={"help": "Min height e.g 128."}) - class VLM_GRPO_DataCollator: - ''' + """ VLM custom data collator for GRPO training, mainly to loads and resize images on the fly. - ''' + """ + def __init__(self, image_dir, min_pixels, max_pixels): self.image_dir = image_dir self.min_pixels = min_pixels self.max_pixels = max_pixels - def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: - + def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: for e in examples: if isinstance(e["image"], str): - item = { + item = { "role": "user", "content": [ {"type": "image", "image": os.path.join("file://" + self.image_dir, e["image"])}, ], } if self.max_pixels is not None: - item["content"][0]["max_pixels"] = self.max_pixels - item["content"][0]["min_pixels"] = self.min_pixels + item["content"][0]["max_pixels"] = self.max_pixels + item["content"][0]["min_pixels"] = self.min_pixels image, _ = process_vision_info([item]) image = image[0] e["image"] = image return examples - + def print_input_config(): - ''' + """ Prints the user provided input config. - ''' + """ args = sys.argv[1:] if "--config" in args and PartialState().is_main_process: - with open(args[args.index("--config")+1], 'r') as file: + with open(args[args.index("--config") + 1]) as file: print("Input config:", file.read().strip()) def vlm_data_format_grpo(sample): - ''' + """ Formats the sample data for GRPO training. - ''' + """ prompt = [{"role": "user", "content": sample["conversations"][0]["value"]}] output = {"id": sample["id"], "solution": sample["solution"], "prompt": prompt, "image": sample["image"]} # print(f"output: {output}") @@ -97,9 +91,9 @@ def vlm_data_format_grpo(sample): def main(script_args, training_args, model_args): - ''' + """ Main function for the GRPO training script. - ''' + """ set_seed(training_args.seed) ############### @@ -126,7 +120,6 @@ def main(script_args, training_args, model_args): logger.info(f"Script parameters {script_args}") logger.info(f"Data parameters {training_args}") - # Check for the last checkpoint if resuming from a previous run last_checkpoint = None if os.path.isdir(training_args.output_dir): @@ -137,23 +130,25 @@ def main(script_args, training_args, model_args): ################ # Load datasets ################ - if PartialState().num_processes > 8: - datasets.disable_progress_bars() + if PartialState().num_processes > 8: + datasets.disable_progress_bars() min_pixels = script_args.min_image_height**2 if script_args.min_image_height is not None else None max_pixels = script_args.max_image_height**2 if script_args.max_image_height is not None else None - my_dataset = datasets.load_dataset("json", data_files=os.path.join(script_args.dataset_root, script_args.dataset_path), split='train', streaming=script_args.dataset_streaming) + my_dataset = datasets.load_dataset( + "json", data_files=os.path.join(script_args.dataset_root, script_args.dataset_path), split="train", streaming=script_args.dataset_streaming + ) train_dataset = my_dataset.map(vlm_data_format_grpo, remove_columns=["conversations"]) if not isinstance(train_dataset, datasets.IterableDataset): logger.info(f"Created dataset mixture with {len(train_dataset)} examples") - training_args.accelerator_config.dispatch_batches = False - if script_args.dataset_streaming: # streaming is not supported with GRPOTrainer yet, but just in case we use it in the future + training_args.accelerator_config.dispatch_batches = False + if script_args.dataset_streaming: # streaming is not supported with GRPOTrainer yet, but just in case we use it in the future train_dataset = train_dataset.shuffle() training_args.dataloader_drop_last = True - training_args.ignore_data_skip = True + training_args.ignore_data_skip = True ############################# # Setup model @@ -164,7 +159,7 @@ def main(script_args, training_args, model_args): torch_dtype=model_args.torch_dtype, trust_remote_code=model_args.trust_remote_code, attn_implementation=model_args.attn_implementation, - use_cache=False + use_cache=False, ) ######################## @@ -175,11 +170,10 @@ def main(script_args, training_args, model_args): "format": format_reward, "tag_count": tag_count_reward, "accuracy_reward_hard": accuracy_reward_hard, - "soft_overshort_punishment": get_soft_overshort_punishment + "soft_overshort_punishment": get_soft_overshort_punishment, } reward_funcs = [REWARD_FUNCS_REGISTRY[func] for func in script_args.reward_funcs] - ############### # GRPO trainer ############### @@ -191,8 +185,7 @@ def main(script_args, training_args, model_args): ) # streaming is not supported for grpo yet, so we load images on the fly in custom data_collator - trainer.data_collator = VLM_GRPO_DataCollator(image_dir=script_args.image_dir, min_pixels=min_pixels, max_pixels=max_pixels) - + trainer.data_collator = VLM_GRPO_DataCollator(image_dir=script_args.image_dir, min_pixels=min_pixels, max_pixels=max_pixels) ############### # Training loop @@ -224,22 +217,23 @@ def main(script_args, training_args, model_args): trainer.model.config.use_cache = True trainer.model.config.save_pretrained(training_args.output_dir) - - logger.info(f"All done! Congrats reaching this point! Please consider citing this work and starring the repo if you found it useful: https://github.com/NVIDIA-Medtech/NV-Reason-CXR") + logger.info( + "All done! Congrats reaching this point! Please consider citing this work and starring the repo if you found it useful: https://github.com/NVIDIA-Medtech/NV-Reason-CXR" + ) if __name__ == "__main__": - ''' + """ Main entry point for the GRPO training script. - ''' + """ print_input_config() - parser = TrlParser((VLMScriptArguments, GRPOConfig, ModelConfig)) # parse arguments + parser = TrlParser((VLMScriptArguments, GRPOConfig, ModelConfig)) # parse arguments script_args, training_args, model_args = parser.parse_args_and_config() # set WANDB run name if training_args.run_name is None: - training_args.run_name = training_args.output_dir.split("/")[-1] + training_args.run_name = training_args.output_dir.split("/")[-1] main(script_args, training_args, model_args) diff --git a/train/vlm_rewards.py b/train/vlm_rewards.py index 0ddff6f..13a24fb 100644 --- a/train/vlm_rewards.py +++ b/train/vlm_rewards.py @@ -1,8 +1,9 @@ import re from dataclasses import dataclass + ## 14 classes based on CheXpert -@dataclass(frozen = True) +@dataclass(frozen=True) class CX: Atelectasis: str = "Atelectasis" Cardiomegaly: str = "Cardiomegaly" @@ -18,8 +19,24 @@ class CX: Pneumonia: str = "Pneumonia" Pneumothorax: str = "Pneumothorax" Support_Devices: str = "Support Devices" + def get_list(): - return [CX.Atelectasis, CX.Cardiomegaly, CX.Consolidation, CX.Edema, CX.Enlarged_Cardiomediastinum, CX.Fracture, CX.Lung_Lesion, CX.Lung_Opacity, CX.No_Finding, CX.Pleural_Effusion, CX.Pleural_Other, CX.Pneumonia, CX.Pneumothorax, CX.Support_Devices] + return [ + CX.Atelectasis, + CX.Cardiomegaly, + CX.Consolidation, + CX.Edema, + CX.Enlarged_Cardiomediastinum, + CX.Fracture, + CX.Lung_Lesion, + CX.Lung_Opacity, + CX.No_Finding, + CX.Pleural_Effusion, + CX.Pleural_Other, + CX.Pneumonia, + CX.Pneumothorax, + CX.Support_Devices, + ] def accuracy_reward(completions, solution, **kwargs): @@ -31,33 +48,30 @@ def accuracy_reward(completions, solution, **kwargs): CX_list = set([e.lower() for e in CX.get_list()]) for content, sol in zip(contents, solution): - # print(f"content: {content} sol: {sol}\n\n") - gold_parsed = set([e.strip().lower() for e in sol.strip().split(",")]) - gold_parsed = gold_parsed & CX_list #filter - + gold_parsed = set([e.strip().lower() for e in sol.strip().split(",")]) + gold_parsed = gold_parsed & CX_list # filter pattern = r"(.*?)" match = re.search(pattern, content, re.DOTALL) if match: answer_parsed = match.group(1) answer_parsed = set([e.strip().lower() for e in answer_parsed.strip().split(",")]) - answer_parsed = answer_parsed & CX_list #filter + answer_parsed = answer_parsed & CX_list # filter intersect = gold_parsed & answer_parsed union = gold_parsed | answer_parsed - reward = float(len(intersect)) / len(union) if len(union) > 0 else 1.0 # if both gold and answer are empty, reward 1.0 (correct) + reward = float(len(intersect)) / len(union) if len(union) > 0 else 1.0 # if both gold and answer are empty, reward 1.0 (correct) else: reward = 0.0 - + rewards.append(reward) return rewards - def format_reward(completions, **kwargs) -> list[float]: """Reward function that checks if the completion has a specific format.""" # print(f"format_reward kwargs: {kwargs} completions: {completions}") @@ -70,20 +84,18 @@ def format_reward(completions, **kwargs) -> list[float]: rewards = [1.0 if match else 0.0 for match in matches] # debug - i=0 + i = 0 for r, c in zip(rewards, completion_contents): if r == 0.0: print(f"format_reward wrong {i} of {rewards} completion_contents {len(c)}: {c}") print("\n\n") - i=i+1 + i = i + 1 return rewards def tag_count_reward(completions, **kwargs) -> list[float]: - """Reward function that checks if we produce the desired number of think and answer tags associated with `format_reward()`. - - """ + """Reward function that checks if we produce the desired number of think and answer tags associated with `format_reward()`.""" def count_tags(text: str) -> float: count = 0.0 @@ -101,31 +113,27 @@ def count_tags(text: str) -> float: return [count_tags(c) for c in contents] - - - def accuracy_reward_hard(completions, solution, **kwargs): """Reward function that checks if the completion is the same as the ground truth.""" # print(f"accuracy_reward kwargs: {kwargs} completions: {completions} solution: {solution}") f_val = format_reward(completions) t_val = tag_count_reward(completions) - t_val = [0.0 if v<1.0 else 1.0 for v in t_val] + t_val = [0.0 if v < 1.0 else 1.0 for v in t_val] - c_val = [f*t for f, t in zip(f_val, t_val)] + c_val = [f * t for f, t in zip(f_val, t_val)] # print(f"c_val: {c_val} for f_val: {f_val} and t_val: {t_val} completions: {completions}") if not any(c_val): - return c_val # return 0.0 for all completions if no format or tag count is correct + return c_val # return 0.0 for all completions if no format or tag count is correct r = accuracy_reward(completions, solution) - rewards = [r*c for r, c in zip(r, c_val)] + rewards = [r * c for r, c in zip(r, c_val)] # print(f"rewards: {rewards} for r before: {r} and c_val: {c_val}") return rewards - def get_soft_overshort_punishment(completion_ids: list[list[int]], **kwargs) -> list[float]: """Reward function that penalizes short completions.""" min_completion_len = 400 @@ -140,5 +148,3 @@ def get_soft_overshort_punishment(completion_ids: list[list[int]], **kwargs) -> rewards.append(max(-1, (completion_length - min_completion_len) / soft_punish_cache)) return rewards - - \ No newline at end of file diff --git a/train/vlm_sft_train.py b/train/vlm_sft_train.py index ec98b57..3f5415c 100644 --- a/train/vlm_sft_train.py +++ b/train/vlm_sft_train.py @@ -2,53 +2,43 @@ import os import sys from dataclasses import dataclass, field -from typing import Dict, Any, List +from typing import Any +import datasets import torch - import transformers -from transformers import set_seed, AutoProcessor, AutoModelForImageTextToText +from datasets import Value +from qwen_vl_utils import process_vision_info +from transformers import AutoModelForImageTextToText, AutoProcessor, Qwen2_5_VLForConditionalGeneration, set_seed from transformers.trainer_utils import get_last_checkpoint +from trl import ModelConfig, ScriptArguments, SFTConfig, SFTTrainer, TrlParser -import datasets -from datasets import Value from accelerate import PartialState -from trl import ( - ModelConfig, - ScriptArguments, - SFTConfig, - SFTTrainer, - TrlParser -) - -from transformers import Qwen2_5_VLForConditionalGeneration -from qwen_vl_utils import process_vision_info - - logger = logging.getLogger(__name__) + @dataclass class VLMScriptArguments(ScriptArguments): - ''' + """ Additional command line arguments for the SFT training script. For a full list of arguments, see the configs/sft_config.yaml file. - ''' + """ dataset_path: str = field(default="grpo.jsonl", metadata={"help": "Path to the dataset json file."}) dataset_root: str = field(default="datalists/", metadata={"help": "Path to the datalists root directory."}) - image_dir: str = field(default="images/", metadata={"help": "Path to the image directory."}) + image_dir: str = field(default="images/", metadata={"help": "Path to the image directory."}) max_image_height: int = field(default=476, metadata={"help": "Max height e.g 476."}) min_image_height: int = field(default=128, metadata={"help": "Min height e.g 128."}) def get_padding_tokens_ids(tokenizer): - ''' + """ Get special tokens ids to mask in the loss computation. - ''' + """ tokenizer = tokenizer.tokenizer if hasattr(tokenizer, "tokenizer") else tokenizer - image_tokens = ["<|image|>", "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>"] + image_tokens = ["<|image|>", "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>"] if hasattr(tokenizer, "image_token"): image_tokens = image_tokens + [tokenizer.image_token] @@ -63,18 +53,18 @@ def get_padding_tokens_ids(tokenizer): class VLM_SFT_DataCollator: - ''' + """ VLM custom data collator for SFT training. - ''' + """ + def __init__(self, processor): self.processor = processor - self.padding_token_ids = get_padding_tokens_ids(processor) # token_ids to ignore in loss computation - - def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: + self.padding_token_ids = get_padding_tokens_ids(processor) # token_ids to ignore in loss computation + def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: ## simplify the structure of the examples examples = [e["messages"] for e in examples] - + ### remove image field if present for messages in examples: for message in messages: @@ -82,47 +72,47 @@ def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: if content["type"] == "image": content.pop("text", None) elif content["type"] == "text": - content.pop("image", None) #remove image field if added by load_dataset - + content.pop("image", None) # remove image field if added by load_dataset # Get the texts and images, and apply the chat template - texts = [self.processor.apply_chat_template(example, tokenize=False) for example in examples] + texts = [self.processor.apply_chat_template(example, tokenize=False) for example in examples] - image_inputs = [process_vision_info(example)[0] for example in examples] + image_inputs = [process_vision_info(example)[0] for example in examples] image_inputs = None if image_inputs[0] is None else image_inputs - batch = self.processor(text=texts, images=image_inputs, return_tensors="pt", padding=True) - - labels = batch["input_ids"].clone() + batch = self.processor(text=texts, images=image_inputs, return_tensors="pt", padding=True) + + labels = batch["input_ids"].clone() labels[torch.isin(labels, self.padding_token_ids)] = -100 # Mask tokens in labels - batch["labels"] = labels + batch["labels"] = labels return batch - + def print_input_config(): - ''' + """ Prints the user provided input config. - ''' + """ args = sys.argv[1:] if "--config" in args and PartialState().is_main_process: - with open(args[args.index("--config")+1], 'r') as file: + with open(args[args.index("--config") + 1]) as file: print("Input config:", file.read().strip()) + def vlm_data_format_dict(sample, image_dir): - ''' + """ Formats the sample data for SFT training. - ''' + """ for message in sample["messages"]: for content in message["content"]: if content["type"] == "image": - content["image"] = os.path.join("file://" + image_dir, content["image"]) #update image path + content["image"] = os.path.join("file://" + image_dir, content["image"]) # update image path content.pop("text", None) elif content["type"] == "text": - content.pop("image", None) #remove image field - - for check_columns in ['id', 'image', 'subject_id', 'study_id', 'solution']: + content.pop("image", None) # remove image field + + for check_columns in ["id", "image", "subject_id", "study_id", "solution"]: if check_columns not in sample: sample[check_columns] = None @@ -130,9 +120,9 @@ def vlm_data_format_dict(sample, image_dir): def main(script_args, training_args, model_args): - ''' + """ Main function for the SFT training script. - ''' + """ set_seed(training_args.seed) @@ -167,23 +157,38 @@ def main(script_args, training_args, model_args): if last_checkpoint is not None and training_args.resume_from_checkpoint is None: logger.info(f"Checkpoint detected, resuming training at {last_checkpoint=}.") - ############### # Setup dataset ############### - features = datasets.Features({"id": Value('string'), "image": Value('string'), "subject_id": Value('string'), "study_id": Value('string'), "solution": Value('string'), "messages": datasets.List({'role': Value('string'), 'content': datasets.List({'type': Value('string'), 'text': Value('string'), 'image': Value('string')})})}) - - my_dataset = datasets.load_dataset("json", data_files=os.path.join(script_args.dataset_root, script_args.dataset_path), split='train', streaming=script_args.dataset_streaming, features=features) + features = datasets.Features( + { + "id": Value("string"), + "image": Value("string"), + "subject_id": Value("string"), + "study_id": Value("string"), + "solution": Value("string"), + "messages": datasets.List( + {"role": Value("string"), "content": datasets.List({"type": Value("string"), "text": Value("string"), "image": Value("string")})} + ), + } + ) + + my_dataset = datasets.load_dataset( + "json", + data_files=os.path.join(script_args.dataset_root, script_args.dataset_path), + split="train", + streaming=script_args.dataset_streaming, + features=features, + ) train_dataset = my_dataset.map(vlm_data_format_dict, fn_kwargs={"image_dir": script_args.image_dir}) - train_dataset = my_dataset.cast(features) + train_dataset = my_dataset.cast(features) - # if streaming + # if streaming if isinstance(train_dataset, datasets.IterableDataset): training_args.dataloader_drop_last = True training_args.accelerator_config.dispatch_batches = False training_args.ignore_data_skip = True - ############### # Setup model ############### @@ -195,12 +200,15 @@ def main(script_args, training_args, model_args): attn_implementation=model_args.attn_implementation, ) - ###################### # Setup Processor ###################### - processor_config={} - if isinstance(model, Qwen2_5_VLForConditionalGeneration) and script_args.min_image_height is not None and script_args.max_image_height is not None: + processor_config = {} + if ( + isinstance(model, Qwen2_5_VLForConditionalGeneration) + and script_args.min_image_height is not None + and script_args.max_image_height is not None + ): processor_config = {"min_pixels": script_args.min_image_height**2, "max_pixels": script_args.max_image_height**2} processor = AutoProcessor.from_pretrained( @@ -208,9 +216,8 @@ def main(script_args, training_args, model_args): use_fast=True, padding_side="right", trust_remote_code=model_args.trust_remote_code, - **processor_config - ) - + **processor_config, + ) ############### # Setup trainer @@ -225,7 +232,6 @@ def main(script_args, training_args, model_args): processing_class=processor, ) - ############### # Training loop ############### @@ -238,7 +244,6 @@ def main(script_args, training_args, model_args): train_result = trainer.train(resume_from_checkpoint=checkpoint) metrics = train_result.metrics - trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() @@ -257,12 +262,12 @@ def main(script_args, training_args, model_args): trainer.model.config.use_cache = True trainer.model.config.save_pretrained(training_args.output_dir) - - logger.info(f"All done! Congrats reaching this point! Please consider citing this work and starring the repo if you found it useful: https://github.com/NVIDIA-Medtech/NV-Reason-CXR") + logger.info( + "All done! Congrats reaching this point! Please consider citing this work and starring the repo if you found it useful: https://github.com/NVIDIA-Medtech/NV-Reason-CXR" + ) if __name__ == "__main__": - print_input_config() # parse arguments @@ -271,6 +276,6 @@ def main(script_args, training_args, model_args): # set WANDB run name if training_args.run_name is None: - training_args.run_name = training_args.output_dir.split("/")[-1] - + training_args.run_name = training_args.output_dir.split("/")[-1] + main(script_args, training_args, model_args)