qwen_multilingual_benchmark.py

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

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

    # Run one specific model on all datasets
    python qwen_multilingual_benchmark.py --model qwen2.5-vl-7b --dataset-dir ./Raw\ CSV/ --output-dir ./results
    
    # Run a model list on all datasets
    python qwen_multilingual_benchmark.py --models qwen2.5-vl-3b qwen2.5-vl-7b --dataset-dir ./Raw\ CSV/
"""

import os
import sys
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

# Import standard libraries
import re
import gc
import json
import time
import argparse
import torch
import pandas as pd
from typing import Dict, List, Optional, Tuple
from tqdm import tqdm

from transformers import (
    pipeline,
    BitsAndBytesConfig
)

# Available Qwen model configuration
QWEN_MODELS = {
    "qwen2-vl-7b": "Qwen/Qwen2-VL-7B-Instruct",
    "qwen2.5-vl-3b": "Qwen/Qwen2.5-VL-3B-Instruct",
    "qwen2.5-vl-7b": "Qwen/Qwen2.5-VL-7B-Instruct",
    "qwen2.5-vl-32b": "Qwen/Qwen2.5-VL-32B-Instruct",
    "qwen2.5-vl-72b": "Qwen/Qwen2.5-VL-72B-Instruct",
    "qwen3-vl-8b": "Qwen/Qwen3-VL-8B-Instruct",
    "qwen3-vl-30b": "Qwen/Qwen3-VL-30B-A3B-Instruct",
    "qwen3-vl-235b": "Qwen/Qwen3-VL-235B-A22B-Instruct",
}

# 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 QwenModelWrapper:
    """Wrapper to manage i Qwen models con pipeline"""
    
    def __init__(self, model_name: str, use_4bit: bool = True):
        self.model_name = model_name
        self.model_id = QWEN_MODELS.get(model_name)
        
        if not self.model_id:
            raise ValueError(f"Model '{model_name}' not available. Qwen models: {list(QWEN_MODELS.keys())}")
        
        # Check whether the model exists locally
        # HuggingFace cache usa: models--<org>--<name>
        org, name = self.model_id.split('/')
        local_model_dir = os.path.join(MODELS_CACHE_DIR, f"models--{org}--{name}", "snapshots")
        
        # If local model is available, use the local path (latest snapshot)
        if os.path.exists(local_model_dir):
            snapshots = os.listdir(local_model_dir)
            if snapshots:
                # Use the latest snapshot available
                latest_snapshot = sorted(snapshots)[-1]
                self.model_path = os.path.join(local_model_dir, latest_snapshot)
                print(f"  ✓ Model found locally: {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.pipeline = 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 with pipeline."""
        print(f"  → Directory models: {MODELS_CACHE_DIR}")
        
        # CONFIGURATION quantizzazione
        model_kwargs = {}
        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"
            )
            model_kwargs["quantization_config"] = quantization_config
            print("  → 4-bit quantization activated")
        
        # Create Qwen pipeline (use local path if available)
        self.pipeline = pipeline(
            "image-text-to-text",
            model=self.model_path,
            model_kwargs=model_kwargs,
            # device=self.device,
            local_files_only=True
        )
        
        print("  ✓ Pipeline initialized")
        print(f"  {get_gpu_memory_info()}")
    
    def generate(self, content: str, image_path: str) -> str:
        """Generate response from the model"""
        try:
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image_path},
                        {"type": "text", "text": content},
                    ],
                },
            ]
            
            outputs = self.pipeline(text=messages, max_new_tokens=100, do_sample=False)
            generated_text = outputs[0]["generated_text"][-1]['content']
            
            return self._extract_answer(generated_text)
        
        except Exception as e:
            print(f"    ✗ Error in generation: {e}")
            return '{"answer": ""}'
    
    def _extract_answer(self, text: str) -> str:
        """Extract the answer in JSON format."""
        # Search for the JSON pattern
        match = re.search(r'{\s*"answer"\s*:\s*"[A-Za-z]"\s*}', text, re.DOTALL)
        if match:
            return match.group(0)
        
        # Try to extract the first valid letter
        if text:
            first_char = text.strip()[0].upper()
            if first_char in ['A', 'B', 'C', 'D', 'E']:
                return f'{{"answer": "{first_char}"}}'
        
        return '{"answer": ""}'
    
    def cleanup(self):
        """Clean up model resources."""
        print(f"\n  ✓ Cleaning up model {self.model_name}...")
        if self.pipeline is not None:
            del self.pipeline
        
        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: QwenModelWrapper,
    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"Numero totale di domande: {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"Colonne mancanti. Richieste: {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)
        # Extract already processed indices (assuming they are saved in results)
        processed_indices = {r.get('Question Index', -1) for r in results}
        print(f"Trovati {len(results)} risultati esistenti - riprendo da dove interrotto")
    
    # 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()
    
    # ⚠️ TEST LIMIT: first 20 questions only
    # MAX_QUESTIONS = 20
    
    for idx, row in tqdm(df.iterrows(), total=len(df), desc=f"Processing {language_code.upper()}"):
        
        # if idx >= MAX_QUESTIONS:
        #     print(f"\n⚠️ Reached limit of {MAX_QUESTIONS} questions - stopping")
        #     break
        
        # Skip if already processed (use Question Index in results to track)
        if idx in processed_indices:
            skipped_count += 1
            continue
        
        question = row['Question']
        
        category = row['Category']
        
        # Build local image path: idx -> (idx+1).jpg con 3 cifre
        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 not found
        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: image not found {image_path}")
                skipped_count += 1
                continue
        
        # Clean 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: 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()
            generated_text = model_wrapper.generate(content, image_path)
            inference_end = time.time()
            inference_time_ms = (inference_end - inference_start) * 1000
            
            response_dict = json.loads(generated_text)
            model_answer = response_dict.get('answer', '').upper()
            is_correct = model_answer == correct_answer_key.upper()
            
            if is_correct:
                correct_count += 1
            total_count += 1
            
            result = {
                'Question Index': idx,
                'Language': language_code.upper(),
                'Category': category,
                'Task': 'Multimodal',
                'Model': model_wrapper.model_id,
                'Question': question,
                'Image Path': image_path,
                'Answer A': answers['AnswerA'],
                'Answer B': answers['AnswerB'],
                'Answer C': answers['AnswerC'],
                'Answer D': answers['AnswerD'],
                'Answer E': answers['AnswerE'],
                'Correct Answer': correct_answer_key.upper(),
                'Model Answer': model_answer,
                'Is Correct': is_correct,
                'Inference Time (ms)': round(inference_time_ms, 2),
            }
            
            results.append(result)
            processed_indices.add(idx)
            
            # Save every 10 questions periodically
            if len(results) % 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:
            error_count += 1
            print(f"\n  ✗ Error: {str(e)[:100]}")
    
    # Final save of results
    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 - len([r for r in results if not r])}") 
    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="Qwen Vision Models - Multilingual Medical Benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Run all Qwen models on all datasets
  python qwen_multilingual_benchmark.py --dataset-dir ./Raw\ CSV/ --output-dir ./results
  
  # Run one specific model
  python qwen_multilingual_benchmark.py --model qwen2.5-vl-7b --dataset-dir ./Raw\ CSV/
  
  # Run a model list
  python qwen_multilingual_benchmark.py --models qwen2.5-vl-3b qwen2.5-vl-7b --dataset-dir ./Raw\ CSV/
        """
    )
    
    parser.add_argument(
        '--dataset-dir',
        type=str,
        default='./Raw CSV/',
        help='Directory containing the CSV files for the datasets (default: ./Raw CSV/)'
    )
    
    parser.add_argument(
        '--images-dir',
        type=str,
        default='./dataset',
        help='Directory containing the images for the datasets (default: ./dataset)'
    )
    
    parser.add_argument(
        '--model',
        type=str,
        choices=list(QWEN_MODELS.keys()),
        help='Run a specific Qwen model'
    )
    
    parser.add_argument(
        '--models',
        type=str,
        nargs='+',
        choices=list(QWEN_MODELS.keys()),
        help='List of models Qwen to run'
    )
    
    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, start fresh'
    )
    
    parser.add_argument(
        '--list-models',
        action='store_true',
        help='Show list of available Qwen 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("\nQwen models available:")
        print("=" * 70)
        for key, value in QWEN_MODELS.items():
            print(f"  {key:20} -> {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"⚠ Warning: Dataset {lang} not found: {csv_path}")
            continue
        
        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(QWEN_MODELS.keys())
        print(f"\n⚠ No model specified - ALL {len(models_to_run)} Qwen models")
        response = input("Continue? [y/N]: ")
        if response.lower() != 'y':
            print("Operation cancelled.")
            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 once for all datasets
            model_wrapper = QwenModelWrapper(model_name, use_4bit=not args.no_4bit)
            
            # Process all datasets with this model
            for dataset_info in datasets_to_process:
                current_combination += 1
                
                print(f"\n{'*'*70}")
                print(f"* Combination {current_combination}/{total_combinations}")
                print(f"* Model: {model_name} | Language: {dataset_info['language'].upper()}")
                print(f"{'*'*70}")
                
                # Genera nome file output
                output_filename = f"{dataset_info['language']}_{model_name}_results.json"
                output_path = os.path.join(args.output_dir, output_filename)
                
                # Process dataset
                result = process_dataset(
                    model_wrapper,
                    dataset_info['path'],
                    dataset_info['code'],
                    output_path,
                    args.images_dir,
                    resume=not args.no_resume
                )
                
                all_results.append(result)
                
                # Pause between datasets of the same model
                if dataset_info != datasets_to_process[-1]:
                    print(f"\n  → Waiting for 3 seconds before the next dataset...")
                    time.sleep(3)
            
            # Model cleanup
            model_wrapper.cleanup()
            
        except Exception as e:
            print(f"\n✗ ERROR in 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['correct']}/{r['total_questions']:>3} {r['time']/60:>11.1f}")
        
        # 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:
            model = r['model']
            if model not in model_stats:
                model_stats[model] = {'accuracies': [], 'times': [], 'total': 0, 'correct': 0}
            model_stats[model]['accuracies'].append(r['accuracy'])
            model_stats[model]['times'].append(r['time'])
            model_stats[model]['total'] += r['total_questions']
            model_stats[model]['correct'] += r['correct']
        
        for model, stats in sorted(model_stats.items(), key=lambda x: sum(x[1]['accuracies'])/len(x[1]['accuracies']), reverse=True):
            avg_accuracy = sum(stats['accuracies']) / len(stats['accuracies'])
            avg_time = sum(stats['times']) / len(stats['times'])
            print(f"{model:<25} Avg Accuracy: {avg_accuracy:>6.2f}% | Avg Time: {avg_time/60:>6.1f}min | Total: {stats['correct']}/{stats['total']}")
        
        # Save overall report
        report_path = os.path.join(args.output_dir, "qwen_multilingual_summary.json")
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump({
                'individual_results': all_results,
                'model_statistics': {
                    model: {
                        'average_accuracy': sum(stats['accuracies']) / len(stats['accuracies']),
                        'average_time_minutes': sum(stats['times']) / len(stats['times']) / 60,
                        'total_questions': stats['total'],
                        'total_correct': stats['correct'],
                        '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()