# The simple, stable solution
model_kwargs = dict(
    torch_dtype=torch.bfloat16,  # Use bfloat16 for its wide numerical range
    device_map="auto",
    attn_implementation="sdpa",
)

​
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, **model_kwargs)

# Install required packages
!pip install --upgrade --quiet transformers datasets evaluate peft trl scikit-learn
​
import os
import re
import torch
import gc
from datasets import load_dataset, ClassLabel
from peft import LoraConfig, PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor
from trl import SFTTrainer, SFTConfig
import evaluate

# Hugging Face authentication using interactive login widget:
from huggingface_hub import notebook_login
notebook_login()

! pip install -q kagglehub
import kagglehub
import os
import pandas as pd
from PIL import Image
from datasets import Dataset, Image as HFImage, Features, ClassLabel
​
# Download the dataset metadata
path = kagglehub.dataset_download("ambarish/breakhis")
print("Path to dataset files:", path)
folds = pd.read_csv('{}/Folds.csv'.format(path))
​
# Filter for 100X magnification from the first fold
folds_100x = folds[folds['mag']==100]
folds_100x = folds_100x[folds_100x['fold']==1]
​
# Get the train/test splits
folds_100x_test = folds_100x[folds_100x.grp=='test']
folds_100x_train = folds_100x[folds_100x.grp=='train']
​
# Define the base path for images
BASE_PATH = "/home/jupyter/.cache/kagglehub/datasets/ambarish/breakhis/versions/4/BreaKHis_v1"

# --- 1. Create Balanced TRAIN Set ---
train_benign_df = folds_100x_train[folds_100x_train['filename'].str.contains('benign')]
train_malignant_df = folds_100x_train[folds_100x_train['filename'].str.contains('malignant')]
min_train_count = min(len(train_benign_df), len(train_malignant_df))
balanced_train_benign = train_benign_df.sample(n=min_train_count, random_state=42)
balanced_train_malignant = train_malignant_df.sample(n=min_train_count, random_state=42)
balanced_train_df = pd.concat([balanced_train_benign, balanced_train_malignant])
​
# --- 2. Create Balanced TEST Set ---
test_benign_df = folds_100x_test[folds_100x_test['filename'].str.contains('benign')]
test_malignant_df = folds_100x_test[folds_100x_test['filename'].str.contains('malignant')]
min_test_count = min(len(test_benign_df), len(test_malignant_df))
balanced_test_benign = test_benign_df.sample(n=min_test_count, random_state=42)
balanced_test_malignant = test_malignant_df.sample(n=min_test_count, random_state=42)
balanced_test_df = pd.concat([balanced_test_benign, balanced_test_malignant])
​
# --- 3. Get the Final Filename Lists ---
train_filenames = balanced_train_df['filename'].values
test_filenames = balanced_test_df['filename'].values
​
print(f"Balanced Train: {len(train_filenames)} files")
print(f"Balanced Test: {len(test_filenames)} files")

CLASS_NAMES = [
    'benign_adenosis', 'benign_fibroadenoma', 'benign_phyllodes_tumor',
    'benign_tubular_adenoma', 'malignant_ductal_carcinoma',
    'malignant_lobular_carcinoma', 'malignant_mucinous_carcinoma',
    'malignant_papillary_carcinoma'
]
​
def get_label_from_filename(filename):
     filename = filename.replace('\\', '/').lower()
     if '/adenosis/' in filename: return 0
     if '/fibroadenoma/' in filename: return 1
     if '/phyllodes_tumor/' in filename: return 2
     if '/tubular_adenoma/' in filename: return 3
     if '/ductal_carcinoma/' in filename: return 4
     if '/lobular_carcinoma/' in filename: return 5
     if '/mucinous_carcinoma/' in filename: return 6
     if '/papillary_carcinoma/' in filename: return 7
     return -1
​
train_data_dict = {
    'image': [os.path.join(BASE_PATH, f) for f in train_filenames],
    'label': [get_label_from_filename(f) for f in train_filenames]
}
test_data_dict = {
    'image': [os.path.join(BASE_PATH, f) for f in test_filenames],
    'label': [get_label_from_filename(f) for f in test_filenames]
}
features = Features({
    'image': HFImage(),
    'label': ClassLabel(names=CLASS_NAMES)
})
train_dataset = Dataset.from_dict(train_data_dict, features=features).cast_column("image", HFImage())
eval_dataset = Dataset.from_dict(test_data_dict, features=features).cast_column("image", HFImage())
​
print(train_dataset)
print(eval_dataset)

# Define the instruction prompt
PROMPT = """Analyze this breast tissue histopathology image and classify it.
​
Classes (0-7):
0: benign_adenosis
1: benign_fibroadenoma
2: benign_phyllodes_tumor
3: benign_tubular_adenoma
4: malignant_ductal_carcinoma
5: malignant_lobular_carcinoma
6: malignant_mucinous_carcinoma
7: malignant_papillary_carcinoma
​
Answer with only the number (0-7):"""
​
def format_data(example):
    """Format examples into the chat-style messages MedGemma expects."""
    example["messages"] = [
        {
            "role": "user",
            "content": [
                {"type": "image"},
                {"type": "text", "text": PROMPT},
            ],
        },
        {
            "role": "assistant",
            "content": [
                {"type": "text", "text": str(example["label"])},
            ],
        },
    ]
    return example
​
# Apply formatting
formatted_train = train_dataset.map(format_data, batched=False)
formatted_eval = eval_dataset.map(format_data, batched=False)
​
print("✓ Data formatted with instruction prompts")

MODEL_ID = "google/medgemma-4b-it"
​
# Model configuration
model_kwargs = dict(
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa",
)
​
model = AutoModelForImageTextToText.from_pretrained(MODEL_ID, **model_kwargs)
processor = AutoProcessor.from_pretrained(MODEL_ID)
​
# Ensure right padding for training
processor.tokenizer.padding_side = "right"

# Helper functions to run evaluation
accuracy_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")
​
def compute_metrics(predictions, references):
    return {
        **accuracy_metric.compute(predictions=predictions, references=references),
        **f1_metric.compute(predictions=predictions, references=references, average="weighted")
    }
​
def postprocess_prediction(text):
    """Extract just the number from the model's text output."""
    digit_match = re.search(r'\b([0-7])\b', text.strip())
    return int(digit_match.group(1)) if digit_match else -1
​
def batch_predict(model, processor, prompts, images, batch_size=8, max_new_tokens=40):
    """A function to run inference in batches."""
    predictions = []
    for i in range(0, len(prompts), batch_size):
        batch_texts = prompts[i:i + batch_size]
        batch_images = [[img] for img in images[i:i + batch_size]]
​
        inputs = processor(text=batch_texts, images=images, padding=True, return_tensors="pt").to("cuda", torch.bfloat16)
        prompt_lengths = inputs["attention_mask"].sum(dim=1)
​
        with torch.inference_mode():
            outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False, pad_token_id=processor.tokenizer.pad_token_id)
​
        for seq, length in zip(outputs, prompt_lengths):
            generated = processor.decode(seq[length:], skip_special_tokens=True)
            predictions.append(postprocess_prediction(generated))
​
    return predictions
​
# Prepare data for evaluation
eval_prompts = [processor.apply_chat_template([msg[0]], add_generation_prompt=True, tokenize=False) for msg in formatted_eval["messages"]]
eval_images = formatted_eval["image"]
eval_labels = formatted_eval["label"]
​
# Run baseline evaluation
print("Running baseline evaluation...")
baseline_preds = batch_predict(model, processor, eval_prompts, eval_images)
baseline_metrics = compute_metrics(baseline_preds, eval_labels)
​
print(f"\n{'BASELINE RESULTS':-^80}")
print(f"Accuracy: {baseline_metrics['accuracy']:.1%}")
print(f"F1 Score: {baseline_metrics['f1']:.3f}")
print("-" * 80)

# LoRA Configuration
peft_config = LoraConfig(
    r=8,
    lora_alpha=16,
    lora_dropout=0.1,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
)
​
# Custom data collator to handle images and text
def collate_fn(examples):
    texts, images = [], []
    for example in examples:
        images.append([example["image"]])
        texts.append(processor.apply_chat_template(example["messages"], add_generation_prompt=False, tokenize=False).strip())
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
    labels = batch["input_ids"].clone()
    labels[labels == processor.tokenizer.pad_token_id] = -100
    image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.tokenizer.special_tokens_map["boi_token"])
    labels[labels == image_token_id] = -100
    labels[labels == 262144] = -100
    batch["labels"] = labels
    return batch
​
# Training arguments
training_args = SFTConfig(
    output_dir="medgemma-breastcancer-finetuned",
    num_train_epochs=5,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=8,
    gradient_checkpointing=True,
    optim="paged_adamw_8bit",
    learning_rate=5e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,  # Warm up LR for first 3% of training
    max_grad_norm=0.3,  # Clip gradients to prevent instability
    bf16=True,  # Use bfloat16 precision
    logging_steps=10,
    save_strategy="steps",
    save_steps=100,
    eval_strategy="epoch",
    push_to_hub=False,
    report_to="none",
    gradient_checkpointing_kwargs={"use_reentrant": False},
    dataset_kwargs={"skip_prepare_dataset": True},
    remove_unused_columns=False,
    label_names=["labels"], 
)
​
# Initialize and run the trainer
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=formatted_train,
    eval_dataset=formatted_eval,
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)
​
print("Starting training...")
trainer.train()
trainer.save_model()

# Clear memory and load the final model
del model
torch.cuda.empty_cache()
gc.collect()
​
# Load base model again
base_model = AutoModelForImageTextToText.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="sdpa"
)
​
# Load LoRA adapters and merge them into a single model
finetuned_model = PeftModel.from_pretrained(base_model, training_args.output_dir)
finetuned_model = finetuned_model.merge_and_unload()
​
# Configure for generation
finetuned_model.generation_config.max_new_tokens = 50
finetuned_model.generation_config.pad_token_id = processor_finetuned.tokenizer.pad_token_id
finetuned_model.config.pad_token_id = processor_finetuned.tokenizer.pad_token_id
​
# Load the processor and run evaluation
processor_finetuned = AutoProcessor.from_pretrained(training_args.output_dir)
finetuned_preds = batch_predict(finetuned_model, processor_finetuned, eval_prompts, eval_images, batch_size=4)
finetuned_metrics = compute_metrics(finetuned_preds, eval_labels)

--- 8-Class Classification (0-7) ---
Model                Accuracy     F1 (Weighted)
-----------------------------------------------
Baseline                  32.6%         0.241
Fine-tuned                87.2%         0.865
-----------------------------------------------
​
--- Binary (Benign/Malignant) Classification ---
Model                Accuracy     F1 (Malignant)
-----------------------------------------------
Baseline                  59.6%         0.639
Fine-tuned                99.0%         0.991
-----------------------------------------------