medgemma-27b-vindr-disease-only-lora

MedGemma 27B LoRA Adapter Model Card

Model Summary

This repository contains a LoRA adapter fine-tuned on top of google/medgemma-27b-it for chest X-ray finding prediction.

The adapter was trained on a processed VinDr-CXR dataset and is intended to predict a list of visible findings from a chest X-ray image using a constrained label set.

This project currently includes three adapter variants:

Base Model Information

• Base model: google/medgemma-27b-it
• Model type: vision-language instruction-tuned generative model
• Adaptation method: QLoRA / LoRA adapter fine-tuning
• Quantization during training: 4-bit NF4 with BitsAndBytes
• Primary task: chest X-ray finding classification as structured text output
• Input modality: chest X-ray image
• Output format: JSON-like finding list

Example output:

{"findings": ["Pleural effusion", "Cardiomegaly"]}

Intended Use

This adapter is designed for:

• research and internal experimentation
• future student handoff and reproducibility

Training Data

• Dataset family: processed VinDr-CXR
• Image type: chest X-ray images
• Supervision target: finding labels stored in dataset metadata
• Training prompt style: one image + instruction prompt + structured assistant answer

The model was trained to output a list of findings only. It was not trained to:

• produce calibrated confidence scores
• explain reasoning
• generate free-form radiology reports
• reliably handle very rare diseases with limited training coverage

Label Space

The training notebook uses the following disease label set:

• Aortic enlargement
• Atelectasis
• Calcification
• Cardiomegaly
• Clavicle fracture
• Consolidation
• Edema
• Emphysema
• Enlarged PA
• ILD
• Infiltration
• Lung Opacity
• Lung cavity
• Lung cyst
• Mediastinal shift
• Nodule/Mass
• Pleural effusion
• Pleural thickening
• Pneumothorax
• Pulmonary fibrosis
• Rib fracture
• Other lesion
• COPD
• Lung tumor
• Pneumonia
• Tuberculosis
• Other diseases

Training Setup

The attached fine-tuning notebook uses the following configuration:

• Seed: 33
• Max sequence length: 2048
• Epochs: 1.0
• Learning rate: 2e-4
• Train batch size per device: 1
• Eval batch size per device: 1
• Gradient accumulation steps: 8
• Warmup ratio: 0.03
• Optimizer: paged_adamw_8bit
• LR scheduler: cosine
• Gradient checkpointing: enabled
• bf16: enabled

LoRA configuration

• r: 16
• alpha: 16
• dropout: 0.05
• target modules: all-linear
• task type: CAUSAL_LM
• bias: none

Quantization configuration

• 4-bit loading enabled
• double quantization enabled
• quant type: nf4
• compute dtype: bfloat16
• quant storage: bfloat16

Prompt Format

The adapter is trained with a system prompt plus a user prompt that asks the model to identify likely findings from the image and return them in a structured format.

The assistant target is serialized as JSON, for example:

{"findings": ["Pneumothorax"]}

Limitations

Current limitations observed in this project:

• performance on rare diseases is weaker because those classes are underrepresented in the training data
• output is limited to finding labels only
• confidence level is not currently
• structured output quality may still vary depending on prompt and image difficulty

Recommended Next Steps

1. Add more rare-disease images to improve performance on underrepresented findings.
2. Fine-tune for confidence prediction so the model can output a confidence value for each predicted finding.
3. Compare the three adapters and their combinations systematically on:
   • disease-heavy evaluation sets
   • healthy-only evaluation sets
   • mixed distributions
   • rare-disease subsets
4. Consider a future version that predicts:
   • structured findings
   • per-finding confidence
   • explain reasoning

---

How to Load the Adapter Locally

1) Install dependencies

pip install -U "transformers>=4.50.0" datasets peft accelerate trl bitsandbytes pillow scikit-learn huggingface_hub

2) Load the base model and processor

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL_ID = "google/medgemma-27b-it"
ADAPTER_PATH = "/path/to/your/adapter_or_checkpoint_directory"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_storage=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)
processor.tokenizer.padding_side = "right"

if processor.tokenizer.pad_token is None:
    processor.tokenizer.pad_token = processor.tokenizer.eos_token

base_model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL_ID,
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="eager",
    device_map="auto",
)

model = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
model.eval()

3) Run inference

from PIL import Image

DISEASE_CLASSES = [
    "Aortic enlargement",
    "Atelectasis",
    "Calcification",
    "Cardiomegaly",
    "Clavicle fracture",
    "Consolidation",
    "Edema",
    "Emphysema",
    "Enlarged PA",
    "ILD",
    "Infiltration",
    "Lung Opacity",
    "Lung cavity",
    "Lung cyst",
    "Mediastinal shift",
    "Nodule/Mass",
    "Pleural effusion",
    "Pleural thickening",
    "Pneumothorax",
    "Pulmonary fibrosis",
    "Rib fracture",
    "Other lesion",
    "COPD",
    "Lung tumor",
    "Pneumonia",
    "Tuberculosis",
    "Other diseases",
]

OPTIONS_TEXT = "\n".join(DISEASE_CLASSES)

USER_PROMPT = (
    "Analyze this X-ray image and identify the most likely findings. "
    "Choose findings only from the options below. "
    "Format your output as {'findings': ['disease 1', 'disease 2', ...]}.\n"
    f"\nOptions:\n{OPTIONS_TEXT}"
)

SYSTEM_PROMPT = """
You are a senior, board-certified radiologist with extensive clinical experience.

Base your analysis strictly on clearly visible, high-confidence visual evidence observed in the provided medical image.
If visual findings are subtle, ambiguous, low-confidence, or potentially attributable to image quality, limited field of view, positioning, or artifacts, explicitly state this.
If bounding boxes are present, each box marks a suspected lesion location. If multiple boxes are shown, review all of them.
Focus primarily on the boxed regions, but do not restrict your assessment to the boxes alone; consider the surrounding areas and the full image context as well.
If no boxes are present, evaluate the entire X-ray image.

Clearly distinguish between:
- Directly observed imaging findings
- Inferred clinical interpretations (which must remain tentative)

Do not provide definitive diagnoses.
Do not label findings as abnormal unless there is clear, high-confidence visual evidence.
If the image or provided context is insufficient to support a conclusion, explicitly state uncertainty.
Do not fabricate findings.
Do not assume clinical history, symptoms, prior studies, or patient context not explicitly provided.
Follow the user-requested output structure strictly.
Do not add any information beyond what is explicitly requested.
Use precise, professional radiologic language.
Limit your response to within 1000 tokens.
""".strip()

image = Image.open("/path/to/image.png").convert("RGB")

messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": SYSTEM_PROMPT}],
    },
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {"type": "text", "text": USER_PROMPT},
        ],
    },
]

prompt_text = processor.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)

inputs = processor(
    text=[prompt_text],
    images=[image],
    return_tensors="pt",
    padding=True,
)

inputs = {k: v.to(model.device) for k, v in inputs.items()}

with torch.no_grad():
    generated = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=False,
    )

output = processor.batch_decode(
    generated[:, inputs["input_ids"].shape[1]:],
    skip_special_tokens=True,
)[0]

print(output.strip())

---

How to Continue Fine-Tuning

Option A: Resume from a Trainer checkpoint

Use this when you want to continue training with optimizer / scheduler / trainer state preserved.

Example checkpoint path:

/medgemma_finetune/checkpoints/<adapter_name>/checkpoint-XXX

In your training script or notebook:

trainer.train(resume_from_checkpoint="./medgemma_finetune/checkpoints/<adapter_name>/checkpoint-XXX")

This is the best option if you want a true continuation of the original run.

Option B: Continue training from the saved adapter weights

Use this when you only have the final saved adapter directory and want to start a new training run initialized from the previously learned adapter.

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
from peft import PeftModel

BASE_MODEL_ID = "google/medgemma-27b-it"
PREVIOUS_ADAPTER_PATH = "/home/ubuntu/medgemma_finetune/checkpoints/<adapter_name>"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_quant_storage=torch.bfloat16,
)

processor = AutoProcessor.from_pretrained(BASE_MODEL_ID)

base_model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL_ID,
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    attn_implementation="eager",
    device_map="auto",
)

model = PeftModel.from_pretrained(base_model, PREVIOUS_ADAPTER_PATH, is_trainable=True)

Then create your Trainer again and call:

trainer.train()

This starts a new run from the learned adapter weights, but it does not restore optimizer state from the earlier run.

Suggested workflow for future groups

• Use checkpoint-* directories when you want an exact resume.
• Use the saved adapter directory when you want to branch into a new experiment.

---