save_evaluate.py

import torch
import json
from datasets import load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    pipeline,
    AutoProcessor,
    Gemma3ForConditionalGeneration, # Updated to use the specific Gemma3 class
)
from accelerate import Accelerator # Added
from tqdm import tqdm

# 1. Initialize Accelerator
accelerator = Accelerator()
NAME = "medgemma_largedino_30ep_unshuffled_normdata_1axis"
suffix = ""
model_id ="[PATH_TO_MODEL]"

# Load Model
model = Gemma3ForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2"
)
model = torch.compile(model)
# Use the Processor, not just the Tokenizer
processor = AutoProcessor.from_pretrained(model_id)

task_name = "report_generation"
dataset = load_dataset("json", data_files=f"[PATH_TO_DATASET]/{task_name}_valid.json", split="train")

def format_vqa_for_trl(data):
    image = data["image"]
    conversation = data["conversations"]
    output = []
    output.append({
        "role": "system",
        "content": [{"type": "text", "text": "You are an expert radiologist."}]
    })
    
    for conv in conversation:
        c_from = conv["from"]
        c_val = conv["value"]
        if c_from == "human":
            content = []
            if "<image>" in c_val:
                content.append({"type": "image", "image": image})
            content.append({
                "type": "text", 
                "text": c_val.replace("<image>","").replace(f"<{task_name}>","").replace("\n", "")
            })
            output.append({"role": "user", "content": content})
        else:
            output.append({
                "role": "assistant", 
                "content": [{"type": "text", "text": c_val}]
            })
    return {"messages": output, "images": [image]}

proprocessed_dataset = dataset.map(
    lambda x: format_vqa_for_trl(x), 
    batched=False, 
    remove_columns=dataset.column_names,
    desc="Formatting VQA data"
)


to_keep = []
# 4. Use accelerator to split the dataset indices across your 4 GPUs
with accelerator.split_between_processes(list(range(len(proprocessed_dataset)))) as indices:
    for i in tqdm(indices):
        test_sample = proprocessed_dataset[i]
        
        # Original logic
        prompt = processor.apply_chat_template(
            test_sample["messages"][:2], 
            add_generation_prompt=True, 
            tokenize=False
            )
        inputs = processor(
            text=prompt, 
            images=test_sample["images"], 
            return_tensors="pt"
            ).to(model.device)
        target_dtype = next(model.model.vison_adaptor.parameters()).dtype
        if "pixel_values" in inputs:
            inputs["pixel_values"] = inputs["pixel_values"].to(target_dtype)
        output_ids = model.generate(
            **inputs,
            max_new_tokens=512,
            do_sample=True,
            temperature=0.2,
            top_p=0.95,
            top_k=50,
            repetition_penalty=1.1,
            use_cache=True,
            )
        
        generated_text = processor.decode(output_ids[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        to_keep.append({
            "question": test_sample['messages'][1]['content'][1]["text"],
            "original": test_sample['messages'][2]['content'][0]["text"],
            "generated": generated_text.strip()
        })

import os

PATH_TO_SAVE = "[PATH_TO_SAVE]"

OUT = os.path.join(PATH_TO_SAVE, NAME+suffix)
os.makedirs(OUT, exist_ok=True)

# 5. Save partial results (one file per GPU)
output_path = os.path.join(OUT, f'to_evaluate_gpu_{accelerator.process_index}.json')
with open(output_path, 'w') as json_file:
    json.dump(to_keep, json_file, indent=4)