Create eval_qwen.py

eval_qwen.py

@@ -0,0 +1,205 @@
+import os
+import logging
+from dataclasses import dataclass
+from typing import List, Dict, Optional, Union
+
+import torch
+from datasets import load_dataset
+import json
+from tqdm import tqdm
+from PIL import Image
+import requests
+from io import BytesIO
+import argparse
+from pathlib import Path
+from enum import Enum
+
+# Import custom modules
+from data import (
+    DatasetType,
+    DatasetConfig,
+    get_dataset_config,
+    get_formatted_instruction,
+    process_response,
+    save_descriptions,
+    load_image_dataset,
+    get_processed_response
+)
+from torch.utils.data import Dataset, DataLoader, DistributedSampler
+import torch.distributed as dist
+from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
+from vllm import LLM, SamplingParams
+
+
+import io
+import base64
+from PIL import Image
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('evaluation.log'),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+INSTRUCTION =  "\n\nYour final answer MUST BE put in \\boxed{}."
+
+def pil_to_base64(image_pil, format="PNG"):
+    buffered = io.BytesIO()
+    image_pil.save(buffered, format=format)
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return img_str
+
+def base64_to_pil(base64_string):
+    img_data = base64.b64decode(base64_string)
+    image_pil = Image.open(io.BytesIO(img_data))
+    return image_pil
+
+class InstanceDataset(Dataset):
+
+    def __init__(self, data):
+        self.data = data
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        item = self.data[index]
+        for k in item:
+            if k == 'options' or k == 'choices':
+                if item[k] == None:
+                    item[k] = ""
+                else:
+                    item[k] = str(item[k])
+        if 'image_url' in item:
+            image_url = item['image_url']
+            image_str = pil_to_base64(image_url)
+            item['image_url'] = image_str
+        instance = {'index': index, 'item': item}
+        return instance
+
+def main():
+    parser = argparse.ArgumentParser(description='Evaluate model on various math datasets')
+    parser.add_argument('--dataset', type=str, choices=['mathvista', 'mathverse', 'mathvision', 'mathvision-mini', 'hallusionbench', 'mmmu-pro-vision', 'we-math', 'math500', 'gpqa', 'dynamath', 'logicvista'],
+                      default='mathvista', help='Dataset to evaluate on')
+    parser.add_argument('--model_path', type=str, help='Path to the model', default="Qwen/Qwen3-VL-2B-Instruct")
+    parser.add_argument('--name', type=str, help='model save name', default="plm")
+    parser.add_argument('--bsz', type=int, help='batch size', default=2)
+
+    args = parser.parse_args()
+    
+    # device = int(os.environ['LOCAL_RANK'])
+    # torch.cuda.set_device(f'cuda:{device}')
+    
+    # Configuration
+    dataset_type = DatasetType(args.dataset)
+    dataset_config = get_dataset_config(dataset_type)
+    
+    output_folder = f"./outputs/{dataset_type.value}_{args.name}"
+    os.makedirs(output_folder, exist_ok=True)
+
+    MODEL_PATH = args.model_path
+    processor = AutoProcessor.from_pretrained(MODEL_PATH)
+    vlm = LLM(MODEL_PATH, limit_mm_per_prompt={"image": 1}, tensor_parallel_size=torch.cuda.device_count())
+    sampling_params = SamplingParams(max_tokens=2048, temperature=0.7, top_p=0.8, top_k=20, repetition_penalty=1.0, presence_penalty=1.5)
+
+    # Load dataset
+    logger.info(f"Loading dataset {dataset_config.name}")
+    data = load_image_dataset(dataset_config)
+    
+    # dist.init_process_group()
+    dataset = InstanceDataset(data)
+    # sampler = DistributedSampler(dataset, shuffle=False)
+    dataloader = DataLoader(dataset, batch_size=args.bsz)
+
+    # Load model
+    # local_rank = int(os.environ['LOCAL_RANK'])
+    # logger.info(f"Loaded model {args.model_path} | local rank: {local_rank}")
+
+    for batch in tqdm(dataloader):
+        
+        indices = batch['index']
+
+        run_input_instances = []
+        run_indices = []
+        run_processed_responses = []
+        run_items = []
+        run_formatted_instructions = []
+
+        for j in range(len(indices)):
+            index = indices[j].item()
+            output_file = os.path.join(output_folder, f'{index}.json')
+            global_item = batch['item']
+            if not os.path.exists(output_file):
+                item = {}
+                for k in global_item:
+                    item[k] = global_item[k][j]
+
+                for k in item:
+                    if len(item[k]) > 0:
+                        if k == 'choices' or k == 'options':
+                            # print(f'item[k]: {item[k]}')
+                            try:
+                                item[k] = eval(item[k])
+                            except:
+                                item[k] = item[k]
+                    if k == 'image_url':
+                        item['image_url'] = base64_to_pil(item['image_url'])
+
+                formatted_instruction = get_formatted_instruction(dataset_type, item)
+                formatted_instruction = formatted_instruction + INSTRUCTION
+                
+                if 'image_url' in item:
+                    message = [{"role": "user", "content": [{"type": "image", "image": ""}, {"type": "text", "text": formatted_instruction}]}]
+                else:
+                    message = [{"role": "user", "content": [{"type": "text", "text": formatted_instruction}]}]
+
+                text = processor.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
+                if 'image_url' in item:
+                    input_instance = {'prompt': text, 'multi_modal_data': {'image': item['image_url']}}
+                else:
+                    input_instance = {'prompt': text}
+                
+                # print(f'input_instance: {input_instance}')
+
+                run_input_instances.append(input_instance)
+                run_indices.append(index)
+            
+                processed_response = get_processed_response(dataset_type, item)
+                # print(f'response: {item["response"]} | processed_response: {processed_response} | choices: {item["choices"]} | ')
+                run_processed_responses.append(processed_response)
+                run_items.append(item)
+                run_formatted_instructions.append(formatted_instruction)
+
+            outputs = vlm.generate(run_input_instances, sampling_params=sampling_params)
+
+            for j in range(len(run_indices)):
+                answer = outputs[j].outputs[0].text
+                processed_response = run_processed_responses[j]
+                item = run_items[j]
+                formatted_instruction = run_formatted_instructions[j]
+                
+                if 'image_url' in item:
+                    del item['image_url']
+
+                description = {
+                    'index': j,
+                    'item': json.dumps(item),
+                    'formatted_instruction': formatted_instruction,
+                    'processed_response': processed_response,
+                    'answer': answer
+                }
+
+                with open(output_file, 'w') as f:
+                    json.dump(description, f, indent = 4)
+
+if __name__ == "__main__":
+    main()
+
+# 
+#  VLLM_WORKER_MULTIPROC_METHOD=spawn  VLLM_DISABLE_COMPILE_CACHE=1 CUDA_VISIBLE_DEVICES=3,4,5,6 python eval_qwen_multi_vllm.py --dataset mathvista --name qwen3_vl_2b_instruct_vllm
+#