medical_models_multilingual_benchmark.py

#!/usr/bin/env python3
"""
Medical Vision-Language Models - Multilingual Medical Benchmark Script
Script to run benchmark tests on multimodal medical models
on a multilingual dataset (English, Italian, Spanish).

Usage:
    # Run all models on all datasets
    python medical_models_multilingual_benchmark.py --dataset-dir ./Raw\ CSV/ --output-dir ./results

    # Run one specific model on all datasets
    python medical_models_multilingual_benchmark.py --model chexagent-8b --dataset-dir ./Raw\ CSV/ --output-dir ./results
    
    # Run a model list on all datasets
    python medical_models_multilingual_benchmark.py --models minicpm-v-2.6 internvl2_5-8b --dataset-dir ./Raw\ CSV/
"""

import argparse
import gc
import json
import os
import re
import sys
import tempfile
import time
from pathlib import Path

# Directory to store downloaded models - SET BEFORE IMPORTING TRANSFORMERS
MODELS_CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
os.makedirs(MODELS_CACHE_DIR, exist_ok=True)

# Set ALL environment variables BEFORE importing transformers
os.environ['HF_HOME'] = MODELS_CACHE_DIR
os.environ['HF_HUB_CACHE'] = MODELS_CACHE_DIR
os.environ['HUGGINGFACE_HUB_CACHE'] = MODELS_CACHE_DIR
os.environ['TRANSFORMERS_CACHE'] = MODELS_CACHE_DIR
os.environ['HF_DATASETS_CACHE'] = MODELS_CACHE_DIR
os.environ['TORCH_HOME'] = MODELS_CACHE_DIR
# Force offline mode to avoid download attempts
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['HF_DATASETS_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'

# Import standard libraries
import pandas as pd
import numpy as np
import torch
from PIL import Image
from typing import Dict, List, Optional, Tuple
from tqdm import tqdm

from transformers import (
    AutoProcessor,
    AutoTokenizer,
    AutoModel,
    AutoModelForCausalLM,
    BitsAndBytesConfig
)

# Models configuration available
MEDICAL_MODELS = {
    "chexagent-3b": "StanfordAIMI/CheXagent-2-3b",
    "chexagent-8b": "StanfordAIMI/CheXagent-8b",
    "minicpm-v-2.6": "openbmb/MiniCPM-V-2_6",
    "minicpm-v-2.6-int4": "openbmb/MiniCPM-V-2_6-int4",
    "internvl2_5-8b": "OpenGVLab/InternVL2_5-8B",
    "llava-med-7b": "microsoft/llava-med-v1.5-mistral-7b",
    "pixtral-12b": "mistral-community/pixtral-12b",
    "maira-2": "microsoft/maira-2",
    "medvlm-r1": "JZPeterPan/MedVLM-R1",
}

# Dataset configuration by language
LANGUAGE_DATASETS = {
    "english": {
        "file": "Raw CSV/Extended_MMMED_English.csv",
        "code": "en"
    },
    "italian": {
        "file": "Raw CSV/Extended_MMMED_Italian.csv",
        "code": "it"
    },
    "spanish": {
        "file": "Raw CSV/Extended_MMMED_Spanish.csv",
        "code": "es"
    }
}

# Multilingual prompt templates
PROMPT_TEMPLATES = {
    "en": """You are a medical student who must answer a multiple-choice test.
Given a medical image and a question related to {category}, choose the correct answer from the options.
Question: {question}
A: {answer_a}
B: {answer_b}
C: {answer_c}
D: {answer_d}
E: {answer_e}
You MUST return an answer EXACTLY in JSON format: {{"answer": "letter"}}.
In ANY CASE, assign a letter equal to the most appropriate option among those provided.
Do not make arguments or reasoning in your response.""",
    
    "it": """Sei uno studente di medicina che deve rispondere a un test a scelta multipla.
Data un'immagine medica e una domanda relativa a {category}, scegli la risposta corretta tra le opzioni.
Domanda: {question}
A: {answer_a}
B: {answer_b}
C: {answer_c}
D: {answer_d}
E: {answer_e}
Devi restituire una risposta ESATTAMENTE in formato JSON: {{"answer": "lettera"}}.
In OGNI CASO, assegna una lettera uguale all'opzione più appropriata tra quelle fornite.
Non fare argomentazioni o ragionamenti nella tua risposta.""",
    
    "es": """Eres un estudiante de medicina que debe responder a un examen de opción múltiple.
Dada una imagen médica y una pregunta relacionada con {category}, elige la respuesta correcta entre las opciones.
Pregunta: {question}
A: {answer_a}
B: {answer_b}
C: {answer_c}
D: {answer_d}
E: {answer_e}
DEBES devolver una respuesta EXACTAMENTE en formato JSON: {{"answer": "letra"}}.
En CUALQUIER CASO, asigna una letra igual a la opción más apropiada entre las proporcionadas.
No hagas argumentos ni razonamientos en tu respuesta."""
}


def clean_text(text):
    """Clean and normalize text."""
    if pd.isna(text):
        return ''
    return str(text).strip().lower().replace('\n', ' ').replace(";", "").replace('"', "")


def clear_memory():
    """Release GPU and CPU memory."""
    gc.collect()
    torch.cuda.empty_cache()
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    print("  ✓ GPU memory released")


def get_gpu_memory_info():
    """Get information about GPU memory usage."""
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        return f"GPU Memory - Allocated: {allocated:.2f}GB, Reserved: {reserved:.2f}GB"
    return "GPU not available"


class MedicalModelWrapper:
    """Wrapper to manage i models medici multimodali"""
    
    def __init__(self, model_name: str, use_4bit: bool = True):
        self.model_name = model_name
        self.model_id = MEDICAL_MODELS.get(model_name)
        
        if not self.model_id:
            raise ValueError(f"Modello '{model_name}' non disponibile. models: {list(MEDICAL_MODELS.keys())}")
        
        # Check whether the model exists locally
        org, name = self.model_id.split('/')
        local_model_dir = os.path.join(MODELS_CACHE_DIR, f"models--{org}--{name}", "snapshots")
        
        if os.path.exists(local_model_dir):
            snapshots = os.listdir(local_model_dir)
            if snapshots:
                latest_snapshot = sorted(snapshots)[-1]
                self.model_path = os.path.join(local_model_dir, latest_snapshot)
                print(f"  → Usando modello locale: {self.model_path}")
            else:
                self.model_path = self.model_id
        else:
            self.model_path = self.model_id
        
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = None
        self.processor = None
        self.tokenizer = None
        
        print(f"\n{'='*70}")
        print(f"Model initialization: {model_name}")
        print(f"HuggingFace ID: {self.model_id}")
        print(f"Device: {self.device}")
        print(f"{'='*70}")
        
        self._initialize(use_4bit)
        
    def _initialize(self, use_4bit: bool):
        """Initialize the model."""
        print(f"  → Directory models: {MODELS_CACHE_DIR}")
        
        # QUANTIZATION CONFIGURATION
        quantization_config = None
        if use_4bit and torch.cuda.is_available():
            quantization_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4"
            )
            print("  → 4-bit quantization activated")
        
        # Model-type specific loading
        if "chexagent" in self.model_name:
            self._init_chexagent(quantization_config)
        elif "minicpm" in self.model_name:
            self._init_minicpm(quantization_config)
        elif "internvl" in self.model_name:
            self._init_internvl(quantization_config)
        elif "llava-med" in self.model_name:
            self._init_llava_med(quantization_config)
        elif "pixtral" in self.model_name:
            self._init_pixtral(quantization_config)
        elif "maira" in self.model_name:
            self._init_maira(quantization_config)
        elif "medvlm" in self.model_name:
            self._init_medvlm(quantization_config)
        
        print(f"  {get_gpu_memory_info()}")
    
    def _init_chexagent(self, quantization_config):
        """Initialize CheXAgent."""
        print("  → Type: CheXAgent")
        
        # Differentiate between CheXAgent-3B and CheXAgent-8B (they have different APIs!)
        is_8b = "8b" in self.model_name.lower()
        
        try:
            if is_8b:
                # CheXAgent-8B uses AutoProcessor
                self.processor = AutoProcessor.from_pretrained(
                    self.model_path,
                    trust_remote_code=True,
                    local_files_only=True
                )
                
                # DO NOT use torch_dtype with quantization for 8B!
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_path,
                    trust_remote_code=True,
                    device_map="auto",
                    quantization_config=quantization_config,
                    low_cpu_mem_usage=True,
                    local_files_only=True
                )
                
                try:
                    self.model.tie_weights()
                except Exception:
                    pass
                    
                self.model.eval()
                print("  ✓ CheXAgent-8B loaded (using AutoProcessor)")
            else:
                # CheXAgent-3B uses AutoTokenizer + from_list_format
                self.tokenizer = AutoTokenizer.from_pretrained(
                    self.model_path, 
                    trust_remote_code=True,
                    local_files_only=True
                )
                if self.tokenizer.pad_token_id is None:
                    self.tokenizer.pad_token = self.tokenizer.eos_token
                
                # For 3B: float16 or bfloat16
                dtype = torch.bfloat16 if (hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported()) else torch.float16
                
                self.model = AutoModelForCausalLM.from_pretrained(
                    self.model_path,
                    device_map="auto",
                    trust_remote_code=True,
                    quantization_config=quantization_config,
                    low_cpu_mem_usage=True,
                    local_files_only=True
                )
                
                # Convert the full model to the correct dtype (critical for 3B!)
                self.model = self.model.to(dtype).eval()
                self.model.model.visual = self.model.model.visual.to(dtype)
                
                print(f"  ✓ CheXAgent-3B loaded (dtype: {dtype})")
            
        except Exception as e:
            error_msg = str(e)
            if "local_files_only" in error_msg.lower() or "network" in error_msg.lower():
                print(f"\n  ⚠ Error: Missing components for offline mode")
                print(f"  → Probably missing: BlipImageProcessor or XraySigLIP")
                print(f"  → Error: {error_msg[:200]}")
                print(f"\n  To resolve, download the complete model first with:")
                print(f"    python download_medical_models.py --models chexagent-8b")
            else:
                if "transformers version" in error_msg.lower() or "transformers==" in error_msg:
                    import transformers
                    current_version = transformers.__version__
                    print(f"\n  ✗ ERROR: CheXAgent requires a specific version of transformers.")
                    print(f"  Current version: {current_version}")
                    print(f"  Required version: transformers==4.40.0 (for chexagent-3b) or >=4.41.2 (for chexagent-8b)")
                    print(f"\n  To resolve:")
                    if "chexagent-2-3b" in self.model_path.lower() or "chexagent-3b" in self.model_name.lower():
                        print(f"    pip install transformers==4.40.0")
                    else:
                        print(f"    pip install 'transformers>=4.41.2'")
                    print(f"\n  Or, use an alternative model like chexagent-8b with the current version.\n")
            raise
    
    def _init_minicpm(self, quantization_config):
        """Initialize MiniCPM-V."""
        print("  → Type: MiniCPM-V")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True, local_files_only=True)
        
        self.model = AutoModel.from_pretrained(
            self.model_path,
            trust_remote_code=True,
            low_cpu_mem_usage=True,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            quantization_config=quantization_config,
            local_files_only=True,
        )
        
        if torch.cuda.is_available():
            self.model = self.model.to("cuda")
        
        self.model.eval()
        print("  ✓ MiniCPM-V loaded")
    
    def _init_internvl(self, quantization_config):
        """Initialize InternVL2.5."""
        print("  → Type: InternVL2.5")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True, local_files_only=True)
        
        dtype = torch.float16
        
        self.model = AutoModel.from_pretrained(
            self.model_path,
            device_map="auto",
            trust_remote_code=True,
            quantization_config=quantization_config,
            torch_dtype=dtype,
            low_cpu_mem_usage=True,
            local_files_only=True,
        ).eval()
        print("  ✓ InternVL2.5 loaded")
    
    def _init_llava_med(self, quantization_config):
        """Initialize LLaVA-Med."""
        print("  → Type: LLaVA-Med")
        # Import LLaVA-Med
        try:
            from llava.model.language_model.llava_mistral import LlavaMistralForCausalLM
        except ImportError:
            print("  ✗ Error: LLaVA-Med requires custom installation")
            print("  → git clone https://github.com/microsoft/LLaVA-Med.git")
            raise
        
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, use_fast=False, local_files_only=True)
        
        # Determine dtype for model loading based on quantization
        if quantization_config is not None:
            # With quantization, use the same compute dtype
            torch_dtype = quantization_config.bnb_4bit_compute_dtype
        else:
            # Without quantization, use float16 for compatibility
            torch_dtype = torch.float16
        
        self.model = LlavaMistralForCausalLM.from_pretrained(
            self.model_path,
            device_map={"": 0},
            quantization_config=quantization_config,
            torch_dtype=torch_dtype,
            low_cpu_mem_usage=True,
            local_files_only=True,
        ).eval()
        
        # Get vision tower and image processor
        vision_tower = self.model.get_vision_tower()
        if isinstance(vision_tower, (list, tuple)):
            vision_tower = vision_tower[0]
        
        if vision_tower is None:
            raise RuntimeError("model.get_vision_tower() ha restituito None")
        
        # If the vision tower is a wrapper that requires loading, do it now
        if hasattr(vision_tower, "is_loaded") and not vision_tower.is_loaded:
            if hasattr(vision_tower, "load_model"):
                vision_tower.load_model()
        
        # Create .vision_tower attribute if not present, by looking for common names
        if not hasattr(vision_tower, "vision_tower"):
            for attr in dir(vision_tower):
                if "vision" in attr.lower() and not attr.startswith("_"):
                    setattr(vision_tower, "vision_tower", getattr(vision_tower, attr))
                    break
        
        if not hasattr(vision_tower, "vision_tower"):
            raise RuntimeError("Vision tower not initialized: missing .vision_tower attribute")
        
        # Move the core vision model to GPU (the actual CLIP model)
        try:
            vision_tower.vision_tower.to(device=self.device, dtype=torch.float16)
        except Exception as e:
            print(f"  ⚠ Warning: impossible to move vision_tower to GPU (it might already be on the device): {e}")
        
        # Save image processor
        self.image_processor = getattr(vision_tower, "image_processor", None)
        if self.image_processor is None:
            raise RuntimeError("vision_tower.image_processor not available")
        
        print("  ✓ LLaVA-Med loaded")
    
    def _init_pixtral(self, quantization_config):
        """Initialize Pixtral."""
        print("  → Type: Pixtral")
        from transformers import LlavaForConditionalGeneration
        
        self.processor = AutoProcessor.from_pretrained(self.model_path, local_files_only=True)
        
        max_memory = {0: "13GiB", "cpu": "48GiB"}
        
        # Determine dtype for model loading based on bfloat16 support
        if torch.cuda.is_available() and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
            dtype = torch.bfloat16
        else:
            dtype = torch.float16
        
        self.model = LlavaForConditionalGeneration.from_pretrained(
            self.model_path,
            quantization_config=quantization_config,
            device_map="auto",
            max_memory=max_memory,
            low_cpu_mem_usage=True,
            torch_dtype=dtype,
            local_files_only=True,
        ).eval()
        
        self.model.config.use_cache = False
        # Store dtype for later conversion
        self._model_dtype = next(self.model.parameters()).dtype
        print("  ✓ Pixtral loaded")
    
    def _init_maira(self, quantization_config):
        """Initialize MAIRA-2."""
        print("  → Type: MAIRA-2")
        # MAIRA-2 not expose `lm_head`: standard 4-bit quantization of
        # transformers/bitsandbytes can fail during the wrapping of the layers.
        # Safe fallback: load MAIRA-2 without quantization.
        maira_quantization_config = quantization_config
        if quantization_config is not None:
            print("  → MAIRA-2: 4-bit not supported in a reliable way, loading without quantization")
            maira_quantization_config = None

        self.processor = AutoProcessor.from_pretrained(
            self.model_path,
            trust_remote_code=True,
            local_files_only=True,
            use_fast=True,
        )

        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_path,
            device_map="auto",
            quantization_config=maira_quantization_config,
            trust_remote_code=True,
            dtype=dtype,
            low_cpu_mem_usage=True,
            local_files_only=True,
        ).eval()
        print("  ✓ MAIRA-2 loaded")
    
    def _init_medvlm(self, quantization_config):
        """Initialize MedVLM-R1."""
        print("  → Type: MedVLM-R1")
        try:
            from transformers import Qwen2VLForConditionalGeneration
        except ImportError:
            print("  ✗ Error: MedVLM-R1 requires qwen-vl-utils")
            print("  → pip install qwen-vl-utils")
            raise
        
        self.processor = AutoProcessor.from_pretrained(self.model_path, max_pixels=28*28*768, local_files_only=True)
        
        if torch.cuda.is_available():
            major, minor = torch.cuda.get_device_capability()
            is_ampere_plus = major >= 8
            torch_dtype = torch.bfloat16 if is_ampere_plus else torch.float16
            # Try flash_attention_2, but fallback to sdpa if not available
            attn_impl = "sdpa"  # Uses sdpa always to avoid dependency on flash-attention
        else:
            torch_dtype = torch.float32
            attn_impl = "eager"
        
        self.model = Qwen2VLForConditionalGeneration.from_pretrained(
            self.model_path,
            device_map="auto",
            torch_dtype=torch_dtype,
            attn_implementation=attn_impl,
            quantization_config=quantization_config,
            low_cpu_mem_usage=True,
            local_files_only=True,
        ).eval()
        print("  ✓ MedVLM-R1 loaded")
    
    def generate(self, content: str, image_path: str) -> str:
        """Generate response from the model"""
        try:
            image = Image.open(image_path).convert("RGB")
            
            if "chexagent" in self.model_name:
                return self._generate_chexagent(content, image_path)
            elif "minicpm" in self.model_name:
                return self._generate_minicpm(content, image)
            elif "internvl" in self.model_name:
                return self._generate_internvl(content, image)
            elif "llava-med" in self.model_name:
                return self._generate_llava_med(content, image)
            elif "pixtral" in self.model_name:
                return self._generate_pixtral(content, image)
            elif "maira" in self.model_name:
                return self._generate_maira(content, image)
            elif "medvlm" in self.model_name:
                return self._generate_medvlm(content, image_path)
            
        except Exception as e:
            print(f"    ✗ Error in generation: {e}")
            return '{"answer": ""}'
    
    def _generate_chexagent(self, content: str, image_path: str) -> str:
        """Generate with CheXAgent"""
        is_8b = "8b" in self.model_name.lower()
        
        try:
            if is_8b:
                # CheXAgent-8B: uses processor, USER/ASSISTANT format
                image = Image.open(image_path).convert("RGB")
                
                # 8B-specific prompt format
                prompt = f" USER: <s>{content} ASSISTANT: <s>"
                
                # Prepare input using processor
                inputs = self.processor(images=[image], text=prompt, return_tensors="pt")
                
                # Move inputs to the correct device with appropriate dtype
                first_device = next(self.model.parameters()).device
                def _move(v):
                    if torch.is_tensor(v):
                        if v.dtype in (torch.float32, torch.float64):
                            return v.to(first_device, dtype=torch.float16)
                        return v.to(first_device)
                    return v
                
                inputs = {k: _move(v) for k, v in inputs.items()}
                
                # Generate
                with torch.no_grad():
                    output_ids = self.model.generate(
                        **inputs,
                        max_new_tokens=32,
                        do_sample=False
                    )[0]
                
                # Decode only the generated part
                output_text = self.processor.tokenizer.decode(output_ids, skip_special_tokens=True)
                assistant_part = output_text.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in output_text else output_text.strip()
                return self._extract_answer(assistant_part)
            else:
                # CheXAgent-3B: uses tokenizer, from_list_format, apply_chat_template
                items = [{"image": image_path}, {"text": content}]
                query = self.tokenizer.from_list_format(items)
                
                # 3B message format
                conv = [
                    {"from": "system", "value": "You are a helpful assistant."},
                    {"from": "human", "value": query}
                ]
                
                # Generate input_ids
                input_ids = self.tokenizer.apply_chat_template(
                    conv, add_generation_prompt=True, return_tensors="pt"
                ).to(self.device)
                
                attention_mask = torch.ones_like(input_ids, dtype=torch.long, device=self.device)
                
                # Generate response
                with torch.no_grad():
                    output_ids = self.model.generate(
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        max_new_tokens=32,
                        do_sample=False,
                        pad_token_id=self.tokenizer.pad_token_id,
                        use_cache=True
                    )[0]
                
                # Decode only the generated part
                output_text = self.tokenizer.decode(output_ids[input_ids.size(1):], skip_special_tokens=True).strip()
                return self._extract_answer(output_text)
        
        except Exception as e:
            print(f"    ✗ Error in generation CheXAgent: {e}")
            import traceback
            traceback.print_exc()
            return '{"answer": ""}'
    
    def _generate_minicpm(self, content: str, image: Image.Image) -> str:
        """Generate with MiniCPM-V"""
        msgs = [{'role': 'user', 'content': [image, content]}]
        
        res = self.model.chat(
            image=None,
            msgs=msgs,
            tokenizer=self.tokenizer,
            sampling=False,
            max_new_tokens=32
        )
        
        return self._extract_answer(res)
    
    def _generate_internvl(self, content: str, image: Image.Image) -> str:
        """Generate with InternVL2.5"""
        # Custom preprocessing for InternVL
        pixel_values = self._preprocess_internvl_image(image)
        
        generation_config = {
            'max_new_tokens': 32,
            'do_sample': False,
        }
        
        try:
            response = self.model.chat(
                self.tokenizer,
                pixel_values,
                content,
                generation_config,
                history=None,
                return_history=False
            )
        except Exception as e:
            # Fallback: if chat fails, try to use generate directly
            print(f"    Warning: model.chat failed ({e}), trying direct generation")
            # Prepare inputs for generation
            inputs = self.tokenizer(content, return_tensors='pt').to(self.device)
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    pixel_values=pixel_values,
                    max_new_tokens=32,
                    do_sample=False
                )
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return self._extract_answer(response)
    
    def _generate_llava_med(self, content: str, image: Image.Image) -> str:
        """Generate with LLaVA-Med using logits scoring"""
        from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
        from llava.mm_utils import process_images, tokenizer_image_token
        from llava.conversation import conv_templates
        
        # Build prompt with image token
        inp = f"{DEFAULT_IMAGE_TOKEN}\n{content}"
        
        # Use appropriate conversation template
        conv_mode = "v1" if "v1" in conv_templates else (
            "vicuna_v1" if "vicuna_v1" in conv_templates else list(conv_templates.keys())[0]
        )
        
        conv = conv_templates[conv_mode].copy()
        conv.append_message(conv.roles[0], inp)
        conv.append_message(conv.roles[1], None)
        prompt = conv.get_prompt()
        
        # Tokenize prompt
        prompt_ids = tokenizer_image_token(
            prompt, self.tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt"
        ).unsqueeze(0).to(self.device)
        
        # Process image
        image_tensor = process_images([image], self.image_processor, self.model.config)
        if isinstance(image_tensor, (list, tuple)):
            image_tensor = torch.stack(image_tensor, dim=0)
        image_tensor = image_tensor.to(device=self.device, dtype=torch.float16)
        
        # Compute score for each letter
        scores = {}
        for letter in "ABCDE":
            # Try variants with and without space
            best_score = -1e9
            for variant in [f" {letter}", f"{letter}"]:
                cand_ids = self.tokenizer.encode(variant, add_special_tokens=False)
                if not cand_ids:
                    continue
                
                # Compute log-likelihood
                cand_tensor = torch.tensor(cand_ids, dtype=prompt_ids.dtype, device=self.device).unsqueeze(0)
                full_ids = torch.cat([prompt_ids, cand_tensor], dim=1)
                
                with torch.no_grad():
                    outputs = self.model(
                        input_ids=full_ids,
                        images=image_tensor,
                        image_sizes=[image.size],
                        use_cache=False,
                    )
                
                logits = outputs.logits
                prompt_len = prompt_ids.shape[1]
                
                log_score = 0.0
                for k, tok in enumerate(cand_ids):
                    pos = (prompt_len - 1) + k
                    lp = torch.log_softmax(logits[0, pos, :], dim=-1)[tok].item()
                    log_score += float(lp)
                
                best_score = max(best_score, log_score)
            
            scores[letter] = best_score
        
        # Select letter with best score
        pred_letter = max(scores, key=scores.get)
        return f'{{"answer": "{pred_letter}"}}'
    
    def _generate_pixtral(self, content: str, image: Image.Image) -> str:
        """Generate with Pixtral"""
        conversation = [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": content}
                ]
            }
        ]
        
        prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=True)
        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
        
        # Convert pixel_values to the model's correct dtype
        if 'pixel_values' in inputs:
            inputs['pixel_values'] = inputs['pixel_values'].to(self._model_dtype)
        
        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=32, do_sample=False)
        
        generated_text = self.processor.decode(outputs[0], skip_special_tokens=True)
        return self._extract_answer(generated_text)
    
    def _generate_maira(self, content: str, image: Image.Image) -> str:
        """Generate with MAIRA-2"""
        messages = [
            {"role": "user", "content": [
                {"type": "image"},
                {"type": "text", "text": content}
            ]}
        ]
        
        prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = self.processor(images=image, text=prompt, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=32, do_sample=False)
        
        generated_text = self.processor.decode(outputs[0], skip_special_tokens=True)
        return self._extract_answer(generated_text)
    
    def _generate_medvlm(self, content: str, image_path: str) -> str:
        """Generate with MedVLM-R1"""
        try:
            from qwen_vl_utils import process_vision_info
        except ImportError:
            return '{"answer": ""}'
        
        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image_path},
                    {"type": "text", "text": content}
                ]
            }
        ]
        
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        image_inputs, video_inputs = process_vision_info(messages)
        
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt"
        ).to(self.device)
        
        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=64, do_sample=False)
        
        generated_ids = [
            output_ids[len(input_ids):] 
            for input_ids, output_ids in zip(inputs.input_ids, outputs)
        ]
        
        output_text = self.processor.batch_decode(
            generated_ids, 
            skip_special_tokens=True, 
            clean_up_tokenization_spaces=False
        )[0]
        
        return self._extract_answer(output_text)
    
    def _preprocess_internvl_image(self, image: Image.Image):
        """InternVL2.5-specific preprocessing with dynamic tiling"""
        from torchvision import transforms
        
        CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
        CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
        
        transform = transforms.Compose([
            transforms.Resize((448, 448), interpolation=transforms.InterpolationMode.BICUBIC),
            transforms.ToTensor(),
            transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
        ])
        
        # Tiling semplificato
        tensor = transform(image).unsqueeze(0)
        return tensor.to(self.device, dtype=torch.float16)
    
    def _extract_answer(self, text: str) -> str:
        """Extract the answer in JSON format."""
        # Cerca pattern JSON
        match = re.search(r'{\s*"answer"\s*:\s*"[A-Za-z]"\s*}', text, re.DOTALL)
        if match:
            return match.group(0)
        
        # Cerca pattern <answer>X</answer>
        match = re.search(r'<answer>\s*([A-E])\s*</answer>', text, re.IGNORECASE)
        if match:
            return '{"answer": "' + match.group(1).upper() + '"}'
        
        # Prova a estrarre la prima lettera valida
        if text:
            match = re.search(r'\b([A-E])\b', text.strip().upper())
            if match:
                return '{"answer": "' + match.group(1) + '"}'
        
        return '{"answer": ""}'
    
    def cleanup(self):
        """Clean up model resources."""
        print(f"\n  Model cleanup {self.model_name}...")
        if self.model is not None:
            del self.model
        if self.processor is not None:
            del self.processor
        if self.tokenizer is not None:
            del self.tokenizer
        
        clear_memory()
        print(f"  {get_gpu_memory_info()}")


def create_prompt(language: str, category: str, question: str, answers: Dict[str, str]) -> str:
    """Create the prompt in the specified language."""
    template = PROMPT_TEMPLATES.get(language, PROMPT_TEMPLATES["en"])
    
    return template.format(
        category=category,
        question=question,
        answer_a=answers['AnswerA'],
        answer_b=answers['AnswerB'],
        answer_c=answers['AnswerC'],
        answer_d=answers['AnswerD'],
        answer_e=answers['AnswerE']
    )


def process_dataset(
    model_wrapper: MedicalModelWrapper,
    csv_path: str,
    language_code: str,
    output_path: str,
    images_dir: str,
    resume: bool = True
) -> Dict:
    """Process a single dataset with the model"""
    
    print(f"\n{'='*70}")
    print(f"PROCESSING DATASET - {language_code.upper()}")
    print(f"{'='*70}")
    print(f"Dataset: {csv_path}")
    print(f"Output: {output_path}")
    
    # Load dataset
    df = pd.read_csv(csv_path)
    print(f"Total questions: {len(df)}")
    
    required_columns = ['Category', 'Question', 'Image url', 'AnswerA', 'AnswerB',
                        'AnswerC', 'AnswerD', 'AnswerE', 'Correct Answer']
    
    if not all(col in df.columns for col in required_columns):
        raise ValueError(f"Columns missing. Required: {required_columns}")
    
    # Load existing results if present
    results = []
    processed_indices = set()
    
    if resume and os.path.exists(output_path):
        with open(output_path, 'r', encoding='utf-8') as f:
            results = json.load(f)
        processed_indices = {r.get('Question Index', -1) for r in results}
        print(f"Found {len(results)} existing results - resuming from where interrupted")
    
    # Statistics
    correct_count = sum(1 for r in results if r.get('Is Correct', False))
    total_count = len(results)
    skipped_count = 0
    error_count = 0
    
    # Process questions
    start_time = time.time()
    
    for idx, row in tqdm(df.iterrows(), total=len(df), desc=f"Processing {language_code.upper()}"):
        
        # Skip if already processed
        if idx in processed_indices:
            skipped_count += 1
            continue
        
        question = row['Question']
        category = row['Category']
        
        # Build local image path
        if idx in [117, 121]:
            image_filename = f"{idx+1:03d}.png"
        else:
            image_filename = f"{idx+1:03d}.jpg"
        image_path = os.path.join(images_dir, image_filename)
        
        # Skip if image does not exist
        if not os.path.exists(image_path):
            image_filename = f"{idx+1:03d}.png"
            image_path = os.path.join(images_dir, image_filename)
            if not os.path.exists(image_path):
                print(f"\n  ⚠ Skipping idx={idx}: image not found")
                skipped_count += 1
                continue
        
        # Prepare answers and correct answer text
        answers = {k: clean_text(row[k]) for k in ['AnswerA', 'AnswerB', 'AnswerC', 'AnswerD', 'AnswerE']}
        answer2key = {v: k[-1] for k, v in answers.items()}
        correct_answer_text = clean_text(row['Correct Answer'])
        
        if correct_answer_text not in answer2key:
            print(f"\n  ⚠ Skipping idx={idx}: correct answer not in options")
            skipped_count += 1
            continue
        
        correct_answer_key = answer2key[correct_answer_text]
        
        # Create prompt in the dataset language
        content = create_prompt(language_code, category, question, answers)
        
        try:
            # Measure inference time
            inference_start = time.time()
            response_text = model_wrapper.generate(content, image_path)
            inference_time = time.time() - inference_start
            
            # Parse model response JSON
            try:
                response_json = json.loads(response_text)
                predicted_key = response_json.get('answer', '').strip().upper()
            except json.JSONDecodeError:
                predicted_key = ''
            
            # Check correctness
            is_correct = (predicted_key == correct_answer_key)
            if is_correct:
                correct_count += 1
            
            total_count += 1
            
            # Save result
            result_entry = {
                'Question Index': int(idx),
                'Category': category,
                'Question': question,
                'Image': image_filename,
                'Correct Answer': correct_answer_key,
                'Model Response': response_text,
                'Predicted Answer': predicted_key,
                'Is Correct': is_correct,
                'Inference Time (s)': round(inference_time, 3)
            }
            results.append(result_entry)
            
            # Incremental save every 10 questions
            if total_count % 10 == 0:
                with open(output_path, 'w', encoding='utf-8') as f:
                    json.dump(results, f, ensure_ascii=False, indent=2)
        
        except Exception as e:
            print(f"\n  ✗ Error idx={idx}: {e}")
            error_count += 1
            result_entry = {
                'Question Index': int(idx),
                'Category': category,
                'Question': question,
                'Image': image_filename,
                'Correct Answer': correct_answer_key,
                'Model Response': f"ERROR: {str(e)}",
                'Predicted Answer': '',
                'Is Correct': False,
                'Inference Time (s)': 0
            }
            results.append(result_entry)
            total_count += 1
    
    # Final save
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    
    elapsed_time = time.time() - start_time
    accuracy = (correct_count / total_count * 100) if total_count > 0 else 0
    
    # Final report
    print(f"\n{'='*70}")
    print(f"RESULTS - {model_wrapper.model_name} - {language_code.upper()}")
    print(f"{'='*70}")
    print(f"Processing time: {elapsed_time/60:.2f} minutes")
    print(f"Total questions in dataset: {len(df)}")
    print(f"Processed questions: {total_count}")
    print(f"Skipped questions: {skipped_count}")
    print(f"Errors: {error_count}")
    print(f"Correct answers: {correct_count}")
    print(f"Accuracy: {accuracy:.2f}%")
    print(f"Results saved: {output_path}")
    print(f"{'='*70}")
    
    return {
        'model': model_wrapper.model_name,
        'language': language_code,
        'total_questions': total_count,
        'correct': correct_count,
        'accuracy': accuracy,
        'time': elapsed_time,
        'output_file': output_path
    }


def main():
    parser = argparse.ArgumentParser(
        description="Medical Vision-Language Models - Multilingual Medical Benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run all models on all datasets
  python medical_models_multilingual_benchmark.py --dataset-dir ./Raw\ CSV/ --output-dir ./results
  
  # Run one specific model
  python medical_models_multilingual_benchmark.py --model chexagent-8b --dataset-dir ./Raw\ CSV/
  
  # Run a model list
  python medical_models_multilingual_benchmark.py --models minicpm-v-2.6 internvl2_5-8b --dataset-dir ./Raw\ CSV/
        """
    )
    
    parser.add_argument(
        '--dataset-dir',
        type=str,
        default='./Raw CSV/',
        help='Directory containing the CSV files of the datasets (default: ./Raw CSV/)'
    )
    
    parser.add_argument(
        '--images-dir',
        type=str,
        default='./dataset',
        help='Directory containing the images of the dataset (default: ./dataset)'
    )
    
    parser.add_argument(
        '--model',
        type=str,
        choices=list(MEDICAL_MODELS.keys()),
        help='Run a specific model'
    )
    
    parser.add_argument(
        '--models',
        type=str,
        nargs='+',
        choices=list(MEDICAL_MODELS.keys()),
        help='List of models to run (space separated)'
    )
    
    parser.add_argument(
        '--output-dir',
        type=str,
        default='./results',
        help='Directory to save results (default: ./results)'
    )
    
    parser.add_argument(
        '--no-4bit',
        action='store_true',
        help='Disable 4-bit quantization'
    )
    
    parser.add_argument(
        '--no-resume',
        action='store_true',
        help='Do not resume from existing results'
    )
    
    parser.add_argument(
        '--list-models',
        action='store_true',
        help='Show list of available models and exit'
    )
    
    parser.add_argument(
        '--languages',
        type=str,
        nargs='+',
        choices=['english', 'italian', 'spanish'],
        default=['english', 'italian', 'spanish'],
        help='Languages of the datasets to process (default: all)'
    )
    
    args = parser.parse_args()
    
    # Models list
    if args.list_models:
        print("\nAvailable medical models:")
        print("=" * 70)
        for key, value in MEDICAL_MODELS.items():
            print(f"  {key:<25} {value}")
        print("=" * 70)
        return
    
    # Check dataset directory
    if not os.path.exists(args.dataset_dir):
        print(f"✗ Error: Dataset directory not found: {args.dataset_dir}")
        sys.exit(1)
    
    # Check dataset for each language
    datasets_to_process = []
    for lang in args.languages:
        csv_filename = LANGUAGE_DATASETS[lang]['file']
        csv_path = os.path.join(args.dataset_dir, csv_filename)
        
        if not os.path.exists(csv_path):
            print(f"✗ Error: Dataset {lang} not found: {csv_path}")
            sys.exit(1)
        
        datasets_to_process.append({
            'language': lang,
            'code': LANGUAGE_DATASETS[lang]['code'],
            'path': csv_path
        })
    
    if not datasets_to_process:
        print("✗ Error: No datasets found!")
        sys.exit(1)
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Determine which models to run
    if args.model:
        models_to_run = [args.model]
    elif args.models:
        models_to_run = args.models
    else:
        models_to_run = list(MEDICAL_MODELS.keys())
        print(f"\n⚠ No model specified - ALL {len(models_to_run)} models")
        response = input("Continue? [y/N]: ")
        if response.lower() != 'y':
            sys.exit(0)
    
    print(f"\n{'='*70}")
    print(f"CONFIGURATION")
    print(f"{'='*70}")
    print(f"Dataset directory: {args.dataset_dir}")
    print(f"Images directory: {args.images_dir}")
    print(f"Output directory: {args.output_dir}")
    print(f"Models to run: {len(models_to_run)}")
    print(f"Datasets to process: {len(datasets_to_process)} ({', '.join([d['language'] for d in datasets_to_process])})")
    print(f"4-bit quantization: {not args.no_4bit}")
    print(f"Resume: {not args.no_resume}")
    print(f"Device: {'CUDA' if torch.cuda.is_available() else 'CPU'}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"{'='*70}")
    
    # Run benchmark for each model on all datasets
    all_results = []
    total_combinations = len(models_to_run) * len(datasets_to_process)
    current_combination = 0
    
    for model_name in models_to_run:
        print(f"\n\n{'#'*70}")
        print(f"# MODEL: {model_name}")
        print(f"{'#'*70}\n")
        
        try:
            # Initialize model
            model_wrapper = MedicalModelWrapper(model_name, use_4bit=not args.no_4bit)
            
            # Process each dataset
            for dataset_info in datasets_to_process:
                current_combination += 1
                print(f"\n[{current_combination}/{total_combinations}] Processing {model_name} on {dataset_info['language']}")
                
                output_filename = f"{dataset_info['language']}_{model_name}_results.json"
                output_path = os.path.join(args.output_dir, output_filename)
                
                result = process_dataset(
                    model_wrapper=model_wrapper,
                    csv_path=dataset_info['path'],
                    language_code=dataset_info['code'],
                    output_path=output_path,
                    images_dir=args.images_dir,
                    resume=not args.no_resume
                )
                
                all_results.append(result)
                
                # Save checkpoint progress
                progress_path = os.path.join(args.output_dir, "medical_progress_checkpoint.json")
                with open(progress_path, 'w', encoding='utf-8') as f:
                    json.dump({
                        'last_model': model_name,
                        'last_language': dataset_info['language'],
                        'completed_combinations': current_combination,
                        'total_combinations': total_combinations,
                        'all_results': all_results
                    }, f, ensure_ascii=False, indent=2)
            
            # Clean up model before loading the next one
            model_wrapper.cleanup()
            
        except Exception as e:
            print(f"\n✗ ERROR with model {model_name}: {e}")
            import traceback
            traceback.print_exc()
            continue
        
        # Wait between models
        if model_name != models_to_run[-1]:
            print(f"\n⚠ Waiting for 10 seconds before the next model...")
            time.sleep(10)
    
    # Final report for all models and datasets
    print(f"\n\n{'#'*70}")
    print(f"# FINAL REPORT - ALL models and datasets")
    print(f"{'#'*70}\n")
    
    if all_results:
        # Sort by model and language
        all_results.sort(key=lambda x: (x['model'], x['language']))
        
        print(f"{'Model':<25} {'Language':>8} {'Accuracy':>10} {'Questions':>10} {'Time (min)':>12}")
        print("=" * 80)
        for r in all_results:
            print(f"{r['model']:<25} {r['language']:>8} {r['accuracy']:>9.2f}% {r['total_questions']:>10} {r['time']/60:>11.2f}")
        
        # Per-model statistics (average across all languages)
        print(f"\n{'='*80}")
        print("AVERAGE PER MODEL (across all languages)")
        print("=" * 80)
        
        model_stats = {}
        for r in all_results:
            if r['model'] not in model_stats:
                model_stats[r['model']] = {'accuracies': [], 'times': [], 'questions': []}
            model_stats[r['model']]['accuracies'].append(r['accuracy'])
            model_stats[r['model']]['times'].append(r['time'])
            model_stats[r['model']]['questions'].append(r['total_questions'])
        
        for model, stats in sorted(model_stats.items(), key=lambda x: sum(x[1]['accuracies'])/len(x[1]['accuracies']), reverse=True):
            avg_acc = sum(stats['accuracies']) / len(stats['accuracies'])
            avg_time = sum(stats['times']) / len(stats['times'])
            print(f"{model:<25} Accuracy media: {avg_acc:>6.2f}%  Tempo medio: {avg_time/60:>6.2f} min")
        
        # Save overall report
        report_path = os.path.join(args.output_dir, "medical_multilingual_summary.json")
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump({
                'summary': {
                    'total_models': len(models_to_run),
                    'total_languages': len(datasets_to_process),
                    'total_combinations': total_combinations,
                    'timestamp': time.strftime('%Y-%m-%d %H:%M:%S')
                },
                'results': all_results,
                'model_averages': {
                    model: {
                        'avg_accuracy': sum(stats['accuracies']) / len(stats['accuracies']),
                        'avg_time': sum(stats['times']) / len(stats['times']),
                        'languages_tested': len(stats['accuracies'])
                    }
                    for model, stats in model_stats.items()
                }
            }, f, ensure_ascii=False, indent=2)
        
        print(f"\n✓ Overall report saved: {report_path}")
    else:
        print("✗ No results available")
    
    print("\n✓ Processing completed!")


if __name__ == "__main__":
    main()