shrishSVaidya/medgemma-1.5-mm-wsi-module3

MedGemma 1.5 - Bone Marrow WSI Pathology Analysis (Module 3)

🔬 Model Overview

This repository contains a PEFT/LoRA adapter fine-tuned on MedGemma 1.5 4B-IT. It is specifically designed to interpret clinical data and pathology reports to provide concise summaries of abnormal plasma cell infiltration, acting as a Whole Slide Imaging (WSI) analysis proxy for Multiple Myeloma patients.

This adapter was developed as part of a edge deployable Agentic AI solution, utilizing the Mixture of Adapters concept. It acts as Module 3, bridging the gap between raw pathology metrics and actionable diagnostic summaries, which are subsequently fed into a RAG-enabled clinical dashboard.

🔗 Associated Code Repository

The complete source code for training this adapter, as well as the full Agenti AI orchestrator pipeline, can be found on GitHub: here

Base Model Dependency

This is an adapter model. It requires the base weights from Google's MedGemma 1.5 4B-IT.

⚠️ License and Terms of Use

• LoRA Adapter Weights: The adapter weights and associated code in this repository are open-sourced under the Apache 2.0 license.
• Base Model: To use this adapter, you must agree to the Google Health AI Developer Foundations Terms of Use to access the underlying MedGemma 1.5 weights.
• Clinical Disclaimer: This model is for educational and research purposes only. It is not a medical device, is not intended for clinical use, and should not be used to diagnose, treat, or offer medical advice for any disease or condition.

💻 How to Use

Because MedGemma is a vision-language architecture, this adapter relies on the vision-capable cross-attention layers. When running this module as a text-only proxy, it is strictly recommended to utilize a dummy image tensor to stabilize the math layers, and load the model in 4-bit NF4 quantization for VRAM efficiency.

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

# 1. Load Base Model in 4-bit NF4
model_id = "google/medgemma-1.5-4b-it"
processor = AutoProcessor.from_pretrained(model_id)

quant_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.float16
)

base_model = AutoModelForImageTextToText.from_pretrained(
    model_id,
    device_map="auto",
    quantization_config=quant_config
)

# 2. Load this LoRA Adapter
model = PeftModel.from_pretrained(base_model, "shrish/medgemma-1.5-mm-wsi-module3")

# 3. Format Prompt and Dummy Image
PROMPT = f"Analyze this 512x512 Bone Marrow Biopsy patch. Does it contain any plasma cells indicative of Multiple Myeloma?"
messages = [{"role": "user", "content": prompt}]
formatted_prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

# Note: If feeding actual WSI images in the future, replace the dummy image with the real PIL Image.
dummy_image = Image.new('RGB', (512, 512), color='black')

inputs = processor(
    text=formatted_prompt, 
    images=dummy_image,
    return_tensors="pt", 
    padding=True
).to(model.device)
inputs.pop("token_type_ids", None)

# 4. Generate
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=300, do_sample=False)

input_length = inputs["input_ids"].shape[1]
print(processor.decode(outputs[0, input_length:], skip_special_tokens=True))