GroundFlow / QWENVL-PRIVATE /analysis /making_base_model_response.py

sync QWENVL-PRIVATE from /opt/tiger/thothvl_pretrain (HF tokens redacted)

070f264 verified 6 months ago

21.2 kB

import sys
import os
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
sys.path.append(project_root)

from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
try:
    from qwen_src.qwen3_vl.modeling_qwen3_vl_batch import Qwen3VLForConditionalGeneration
except Exception:
    Qwen3VLForConditionalGeneration = None
import torch
import matplotlib.pyplot as plt
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor
import math
from PIL import Image, ImageDraw, ImageFont
import numpy as np


import json
import os
from tqdm import tqdm

import pickle
import argparse
import multiprocessing as mp
import pandas as pd  # <-- Make sure to import pandas
from matching_utils import TextVQAAccuracyEvaluator
from check_gate_label import plot_single_sample


def infer_model_type(model_path, model_type):
    if model_type and model_type != "auto":
        return model_type
    lower = model_path.lower()
    if "qwen3" in lower:
        return "qwen3"
    if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
        return "qwen2_5"
    if "qwen2" in lower:
        return "qwen2"
    return "qwen2_5"


def resolve_processor_path(model_path, model_type, processor_path=None):
    if processor_path:
        return processor_path
    lower = model_path.lower()
    public_prefix = "/root/public/Qwen"

    def to_repo_path(name):
        if "/root/public" in model_path:
            return f"{public_prefix}/{name}"
        return f"Qwen/{name}"

    if model_type == "qwen3":
        if "8b" in lower:
            return to_repo_path("Qwen3-VL-8B-Instruct")
        if "4b" in lower:
            return to_repo_path("Qwen3-VL-4B-Instruct")
    if model_type == "qwen2_5":
        if "7b" in lower:
            return to_repo_path("Qwen2.5-VL-7B-Instruct")
        if "3b" in lower:
            return to_repo_path("Qwen2.5-VL-3B-Instruct")
    if model_type == "qwen2" and "7b" in lower:
        return to_repo_path("Qwen2-VL-7B-Instruct")

    return model_path


def resolve_model_class(model_type):
    if model_type == "qwen3":
        if Qwen3VLForConditionalGeneration is None:
            raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
        return Qwen3VLForConditionalGeneration
    if model_type == "qwen2":
        return Qwen2VLForConditionalGeneration
    return Qwen2_5_VLForConditionalGeneration


def set_twig_flags(model, model_type, two_stage_flag):
    if model_type == "qwen3":
        if hasattr(model, "enable_twig"):
            model.enable_twig = two_stage_flag
        if hasattr(model, "roi_enable2stage"):
            model.roi_enable2stage = two_stage_flag
        if hasattr(model, "language_model"):
            if hasattr(model.language_model, "enable_twig"):
                model.language_model.enable_twig = two_stage_flag
            if hasattr(model.language_model, "roi_enable2stage"):
                model.language_model.roi_enable2stage = two_stage_flag
        return

    if hasattr(model, "model"):
        if hasattr(model.model, "enable_twig"):
            model.model.enable_twig = two_stage_flag
        if hasattr(model.model, "roi_enable2stage"):
            model.model.roi_enable2stage = two_stage_flag
        if hasattr(model.model, "two_stage_vanilla"):
            model.model.two_stage_vanilla = two_stage_flag
    if hasattr(model, "roi_enable2stage"):
        model.roi_enable2stage = two_stage_flag
def get_KL_divergence(logits_list,):
    import torch.nn as nn
    from torch.nn import functional as F
    KL_div = nn.KLDivLoss(reduction='batchmean')
    KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
    return KL_div_list


from torch.utils.data import Dataset, DataLoader


def worker_wrapper(chunk, gpu_id, args, return_dict):
    """
    Wraps the main worker function to handle return values via shared dictionary.
    """

    result = worker_process(chunk, gpu_id, args)
    # Store result in the shared dictionary using gpu_id as key
    return_dict[gpu_id] = result



class QwenEvalDataset(Dataset):
    def __init__(self, questions_chunk, args):
        self.data = questions_chunk
        self.args = args
        self.image_folder = args.image_folder
        self.use_short_prompt = getattr(args, 'use_short_prompt', False)


    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        line = self.data[idx]
        image_file_str = ''
        ground_truth = None
        qs_text = ""
        image = None

        try:

            qs_text = line["text"]
            qs_text = qs_text
            image_file = line["image"]
            image_path = os.path.join(self.image_folder, image_file)
            image = Image.open(image_path).convert('RGB')
            image_file_str = image_file
            ground_truth = line.get("answer") if "answer" in line else line.get("answers", None)

            if "<image>" in qs_text:
                qs_text = qs_text.replace("<image>", "")  # Remove it from the text body

            qs = "<image>\n" + qs_text

            if self.use_short_prompt:
                qs = qs + "\nAnswer the question using a single word or phrase."

            # Return a dict containing everything needed for inference
            return {
                'image': image,
                'prompt_text': qs,
                'meta': {
                    'image_id': image_file_str,
                    'prompt': qs_text,
                    'ground_truth': ground_truth,
                    'line_number': line['original_line_number'],
                    'id': line.get('question_id', line.get('original_line_number', 'unknown'))
                },
                'valid': True
            }

        except Exception as e:
            print(f"Error loading item {idx}: {e}")
            # Return dummy invalid item
            return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}


class SingleResCollator:
    def __init__(self, processor_path, patch_size):
        self.processor_path = processor_path
        self.processor = None
        # Define the requested limits
        self.min_pixels = 128 * patch_size * patch_size
        self.max_pixels = 576 * patch_size * patch_size

    def __call__(self, batch):
        if self.processor is None:
            # Initialize processor once per worker
            self.processor = AutoProcessor.from_pretrained(
                self.processor_path,
                trust_remote_code=True,
                padding_side='left',
                use_fast=True
            )
            # Set the resolution constraints globally for this processor
            self.processor.image_processor.min_pixels = self.min_pixels
            self.processor.image_processor.max_pixels = self.max_pixels

        batch = [item for item in batch if item.get('valid', False)]
        if not batch:
            return None

        # 1. Extract raw data
        raw_images = [item['image'] for item in batch]
        raw_prompts = [item['prompt_text'] for item in batch]
        metas = [item['meta'] for item in batch]

        # 2. Prepare text and images
        current_texts = []
        for prompt_str, img in zip(raw_prompts, raw_images):
            # Qwen-VL-Chat template construction
            messages = [{"role": "user", "content": [{"type": "text", "text": prompt_str}, {"image": img}]}]
            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

            # Specific replacement logic for Qwen2.5-VL
            text = text.replace('<|vision_start|><|image_pad|><|vision_end|>', '') \
                .replace('<image>', '<image><|vision_start|><|image_pad|><|vision_end|>')
            current_texts.append(text)

        try:
            # 3. Process Batch (Single Resolution Setting)
            inputs = self.processor(
                text=current_texts,
                images=raw_images,
                padding=True,
                return_tensors="pt"
            )
            # Attach extra data needed later if necessary (or just return inputs)
            inputs.data['src_images'] = raw_images

            return {
                'metas': metas,
                'inputs': inputs
            }

        except Exception as e:
            print(f"Error processing batch: {e}")
            return None


def worker_process(questions_chunk, gpu_id, args):
    device = f'cuda:{gpu_id}'
    answer_analysis_dict = {}  # Renamed from ms_answer_analysis_dict
    max_new_tokens = 256

    # --- SETUP DATALOADER ---
    batch_size = getattr(args, 'batch_size', 8)
    num_workers = getattr(args, 'num_workers', 4)

    dataset = QwenEvalDataset(questions_chunk, args)

    # Initialize the NEW single resolution collator
    collator = SingleResCollator(args.processor_path, args.patch_size)

    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collator,
        pin_memory=True
    )

    if len(dataset) == 0: return answer_analysis_dict

    try:
        # Load Model
        print(f"Worker {gpu_id}: Loading model...")
        model_class = resolve_model_class(args.model_type)
        model = model_class.from_pretrained(
            args.model_path,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map=device
        )
        model.eval()

        # Flags (keeping your original config)
        two_stage_flag = False
        set_twig_flags(model, args.model_type, two_stage_flag)

        # Tokenizer for decoding
        tokenizer = AutoProcessor.from_pretrained(args.processor_path, trust_remote_code=True).tokenizer

        print(f"Worker {gpu_id}: Starting inference...")

        # --- MAIN INFERENCE LOOP ---
        progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)

        for batch_data in progress_bar:
            if batch_data is None: continue

            metas = batch_data['metas']
            inputs = batch_data['inputs']  # Now getting direct inputs, not inputs_by_res

            try:
                # Move to GPU
                inputs = inputs.to(device)

                with torch.no_grad():
                    output_ids = model.generate(
                        **inputs,
                        max_new_tokens=max_new_tokens,
                        output_attentions=False,
                        pad_token_id=tokenizer.eos_token_id
                    )

                # Decode
                input_ids_len = inputs.input_ids.shape[1]
                generated_ids = output_ids[:, input_ids_len:]
                batch_answers = tokenizer.batch_decode(
                    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
                )

                # --- SAVE RESULTS ---
                for idx, (meta, answer) in enumerate(zip(metas, batch_answers)):
                    line_number = meta['line_number']

                    tmp_sample = {
                        'image': meta['image_id'],
                        'prompt': meta['prompt'],
                        'text': meta['ground_truth'],
                        'id': meta['id'],
                        'response': answer  # Storing single response
                    }
                    answer_analysis_dict[line_number] = tmp_sample

                # Cleanup GPU memory for this batch
                del inputs, output_ids, generated_ids

            except RuntimeError as e:
                msg = "ERROR_OOM" if "out of memory" in str(e) else f"ERROR_RUNTIME: {e}"
                print(f"Worker {gpu_id}: {msg}")
                torch.cuda.empty_cache()
                # Log errors for this batch
                for meta in metas:
                    line_number = meta['line_number']
                    answer_analysis_dict[line_number] = {
                        'id': meta['id'],
                        'response': msg
                    }

    except Exception as e:
        print(f"Worker {gpu_id}: Critical Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        torch.cuda.empty_cache()

    return answer_analysis_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
    parser.add_argument('--model-path', type=str, required=True,
                        help="Path to the model and processor")
    parser.add_argument('--model-type', type=str, default="auto",
                        choices=["auto", "qwen2_5", "qwen2", "qwen3"],
                        help="Model family for loading (auto detects from model path).")
    parser.add_argument('--processor-path', type=str, default="",
                        help="Override processor path if auto detection is incorrect.")
    parser.add_argument('--question-file', type=str, default=None,
                        help="Path to the JSONL question file (E.g., TextVQA)")
    parser.add_argument('--image-folder', type=str, default=None,
                        help="Path to the image directory (for JSONL file)")
    parser.add_argument('--dst-save-path', type=str, required=True,
                        help="Path to save the output .pkl file")
    parser.add_argument('--gpu-ids', type=str, default='0,1',
                        help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
    # --- ADDED: New HF Dataset args ---
    parser.add_argument('--hf_dataset', type=str, default=None,
                        help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
    parser.add_argument('--dataset_split', type=str, default='test',
                        help="Split to use from the Hugging Face dataset (e.g., 'test')")
    parser.add_argument('--is_check_pred', type=bool, default=False,
                        help="Whether to check high-res predictions")
    ### --- END ADDITION ---
    parser.add_argument('--use-short-prompt', action='store_true',
                        help="Whether to use short prompt format")
    args = parser.parse_args()
    args.model_type = infer_model_type(args.model_path, args.model_type)
    args.processor_path = resolve_processor_path(
        args.model_path, args.model_type, args.processor_path
    )
    args.patch_size = 32 if args.model_type == "qwen3" else 28


    # --- 1. Define a Wrapper Function ---
    # This helper must be defined at the top level (global scope) so it can be pickled.
    # --- 2. The Main Execution Block ---
    # Ensure start method is spawn (Crucial for PyTorch+CUDA)
    try:
        mp.set_start_method('spawn', force=True)
    except RuntimeError:
        pass

    # 1. Prepare Data and Split into Chunks
    gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
    num_gpus = len(gpu_ids)
    questions = []

    if args.hf_dataset:
        print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
        from datasets import load_dataset

        dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
        print("Casting 'image' column to load from cache...")
        questions_list = list(dataset)

        new_questions = []
        for i, tmp_ques in enumerate(questions_list):
            tmp_ques['original_line_number'] = i
            new_questions.append(tmp_ques)
        questions = new_questions
        print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")

    elif args.question_file:
        print("Loading and preprocessing questions from JSONL file...")
        if not args.image_folder:
            raise ValueError("--image-folder is required when using --question-file")

        questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]

        new_questions = []
        for i, tmp_ques in enumerate(questions_list):
            qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
            if 'Reference OCR' in qs:
                qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
            tmp_ques['text'] = qs
            tmp_ques['original_line_number'] = i
            new_questions.append(tmp_ques)
        questions = new_questions
        print(f"Loaded {len(questions)} preprocessed questions from local file.")

    else:
        raise ValueError("You must provide either --question-file or --hf_dataset")

    # Create chunks
    chunk_size = math.ceil(len(questions) / num_gpus)
    chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]

    # 2. Run processes in parallel using mp.Process (Manual Management)
    print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")

    # Create a Manager to hold the results safely across processes
    manager = mp.Manager()
    return_dict = manager.dict()
    processes = []

    for i in range(num_gpus):
        chunk_data = chunks[i] if i < len(chunks) else []

        # We pass the shared 'return_dict' to the wrapper
        p = mp.Process(
            target=worker_wrapper,
            args=(chunk_data, gpu_ids[i], args, return_dict)
        )
        p.start()
        processes.append(p)

    # Wait for all processes to finish
    for p in processes:
        p.join()

    # 3. Gather and Save Results
    print("All chunks processed. Gathering results...")
    final_dict = {}

    # Extract results from the shared return_dict
    for gpu_id, sub_dict in return_dict.items():
        final_dict.update(sub_dict)

    #check save directory
    os.makedirs(os.path.dirname(os.path.expanduser(args.dst_save_path)), exist_ok=True)
    with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
        pickle.dump(final_dict, f)

    print(f"✅ Successfully processed and saved {len(final_dict)} answers to {args.dst_save_path}.")


'''
CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
    --question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
    --image-folder ../llava/playground/data/eval/textvqa/train_images \
    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B.pkl \
    --gpu-ids 0,1
CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
    --model-path Qwen/Qwen2.5-VL-7B-Instruct \
    --question-file /home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl \
    --image-folder /home/yuheng/datasets \
    --dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B.pkl \
    --gpu-ids 0,1
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/textvqa/converted_llava_style_train.jsonl \
    --image-folder /root/epfs/sproject/datasets/textvqa/train_images \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/chartqa_qwen2.5vl-7B-short-prompt.pkl \\
    --use-short-prompt \
    --gpu-ids 0,1
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/visual_cot_infovqa_subset15k.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_infovqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
    --image-folder /root/epfs/sproject/datasets \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_docvqa_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
    
python analysis/making_base_model_response.py \
    --model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
    --question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
    --image-folder /root/epfs/sproject/datasets/ \
    --dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-short-prompt.pkl \
    --use-short-prompt \
    --gpu-ids 0,1
'''

	import sys
	import os
	script_dir = os.path.dirname(os.path.abspath(__file__))
	project_root = os.path.dirname(script_dir)
	sys.path.append(project_root)

	from qwen_src.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration
	from qwen_src.qwen2_vl import Qwen2VLForConditionalGeneration
	try:
	from qwen_src.qwen3_vl.modeling_qwen3_vl_batch import Qwen3VLForConditionalGeneration
	except Exception:
	Qwen3VLForConditionalGeneration = None
	import torch
	import matplotlib.pyplot as plt
	from qwen_vl_utils import process_vision_info
	from transformers import AutoProcessor
	import math
	from PIL import Image, ImageDraw, ImageFont
	import numpy as np


	import json
	import os
	from tqdm import tqdm

	import pickle
	import argparse
	import multiprocessing as mp
	import pandas as pd # <-- Make sure to import pandas
	from matching_utils import TextVQAAccuracyEvaluator
	from check_gate_label import plot_single_sample


	def infer_model_type(model_path, model_type):
	if model_type and model_type != "auto":
	return model_type
	lower = model_path.lower()
	if "qwen3" in lower:
	return "qwen3"
	if "qwen2.5" in lower or "qwen2_5" in lower or "qwen2-5" in lower:
	return "qwen2_5"
	if "qwen2" in lower:
	return "qwen2"
	return "qwen2_5"


	def resolve_processor_path(model_path, model_type, processor_path=None):
	if processor_path:
	return processor_path
	lower = model_path.lower()
	public_prefix = "/root/public/Qwen"

	def to_repo_path(name):
	if "/root/public" in model_path:
	return f"{public_prefix}/{name}"
	return f"Qwen/{name}"

	if model_type == "qwen3":
	if "8b" in lower:
	return to_repo_path("Qwen3-VL-8B-Instruct")
	if "4b" in lower:
	return to_repo_path("Qwen3-VL-4B-Instruct")
	if model_type == "qwen2_5":
	if "7b" in lower:
	return to_repo_path("Qwen2.5-VL-7B-Instruct")
	if "3b" in lower:
	return to_repo_path("Qwen2.5-VL-3B-Instruct")
	if model_type == "qwen2" and "7b" in lower:
	return to_repo_path("Qwen2-VL-7B-Instruct")

	return model_path


	def resolve_model_class(model_type):
	if model_type == "qwen3":
	if Qwen3VLForConditionalGeneration is None:
	raise RuntimeError("Qwen3VLForConditionalGeneration not available.")
	return Qwen3VLForConditionalGeneration
	if model_type == "qwen2":
	return Qwen2VLForConditionalGeneration
	return Qwen2_5_VLForConditionalGeneration


	def set_twig_flags(model, model_type, two_stage_flag):
	if model_type == "qwen3":
	if hasattr(model, "enable_twig"):
	model.enable_twig = two_stage_flag
	if hasattr(model, "roi_enable2stage"):
	model.roi_enable2stage = two_stage_flag
	if hasattr(model, "language_model"):
	if hasattr(model.language_model, "enable_twig"):
	model.language_model.enable_twig = two_stage_flag
	if hasattr(model.language_model, "roi_enable2stage"):
	model.language_model.roi_enable2stage = two_stage_flag
	return

	if hasattr(model, "model"):
	if hasattr(model.model, "enable_twig"):
	model.model.enable_twig = two_stage_flag
	if hasattr(model.model, "roi_enable2stage"):
	model.model.roi_enable2stage = two_stage_flag
	if hasattr(model.model, "two_stage_vanilla"):
	model.model.two_stage_vanilla = two_stage_flag
	if hasattr(model, "roi_enable2stage"):
	model.roi_enable2stage = two_stage_flag
	def get_KL_divergence(logits_list,):
	import torch.nn as nn
	from torch.nn import functional as F
	KL_div = nn.KLDivLoss(reduction='batchmean')
	KL_div_list = [KL_div(F.log_softmax(logits_list[-1], dim=1), F.softmax(logits_list[i], dim=1)).item() for i in range(len(logits_list))]
	return KL_div_list


	from torch.utils.data import Dataset, DataLoader


	def worker_wrapper(chunk, gpu_id, args, return_dict):
	"""
	Wraps the main worker function to handle return values via shared dictionary.
	"""

	result = worker_process(chunk, gpu_id, args)
	# Store result in the shared dictionary using gpu_id as key
	return_dict[gpu_id] = result



	class QwenEvalDataset(Dataset):
	def __init__(self, questions_chunk, args):
	self.data = questions_chunk
	self.args = args
	self.image_folder = args.image_folder
	self.use_short_prompt = getattr(args, 'use_short_prompt', False)


	def __len__(self):
	return len(self.data)

	def __getitem__(self, idx):
	line = self.data[idx]
	image_file_str = ''
	ground_truth = None
	qs_text = ""
	image = None

	try:

	qs_text = line["text"]
	qs_text = qs_text
	image_file = line["image"]
	image_path = os.path.join(self.image_folder, image_file)
	image = Image.open(image_path).convert('RGB')
	image_file_str = image_file
	ground_truth = line.get("answer") if "answer" in line else line.get("answers", None)

	if "<image>" in qs_text:
	qs_text = qs_text.replace("<image>", "") # Remove it from the text body

	qs = "<image>\n" + qs_text

	if self.use_short_prompt:
	qs = qs + "\nAnswer the question using a single word or phrase."

	# Return a dict containing everything needed for inference
	return {
	'image': image,
	'prompt_text': qs,
	'meta': {
	'image_id': image_file_str,
	'prompt': qs_text,
	'ground_truth': ground_truth,
	'line_number': line['original_line_number'],
	'id': line.get('question_id', line.get('original_line_number', 'unknown'))
	},
	'valid': True
	}

	except Exception as e:
	print(f"Error loading item {idx}: {e}")
	# Return dummy invalid item
	return {'valid': False, 'meta': {'line_number': line.get('original_line_number', -1)}}


	class SingleResCollator:
	def __init__(self, processor_path, patch_size):
	self.processor_path = processor_path
	self.processor = None
	# Define the requested limits
	self.min_pixels = 128 * patch_size * patch_size
	self.max_pixels = 576 * patch_size * patch_size

	def __call__(self, batch):
	if self.processor is None:
	# Initialize processor once per worker
	self.processor = AutoProcessor.from_pretrained(
	self.processor_path,
	trust_remote_code=True,
	padding_side='left',
	use_fast=True
	)
	# Set the resolution constraints globally for this processor
	self.processor.image_processor.min_pixels = self.min_pixels
	self.processor.image_processor.max_pixels = self.max_pixels

	batch = [item for item in batch if item.get('valid', False)]
	if not batch:
	return None

	# 1. Extract raw data
	raw_images = [item['image'] for item in batch]
	raw_prompts = [item['prompt_text'] for item in batch]
	metas = [item['meta'] for item in batch]

	# 2. Prepare text and images
	current_texts = []
	for prompt_str, img in zip(raw_prompts, raw_images):
	# Qwen-VL-Chat template construction
	messages = [{"role": "user", "content": [{"type": "text", "text": prompt_str}, {"image": img}]}]
	text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

	# Specific replacement logic for Qwen2.5-VL
	text = text.replace('<\|vision_start\|><\|image_pad\|><\|vision_end\|>', '') \
	.replace('<image>', '<image><\|vision_start\|><\|image_pad\|><\|vision_end\|>')
	current_texts.append(text)

	try:
	# 3. Process Batch (Single Resolution Setting)
	inputs = self.processor(
	text=current_texts,
	images=raw_images,
	padding=True,
	return_tensors="pt"
	)
	# Attach extra data needed later if necessary (or just return inputs)
	inputs.data['src_images'] = raw_images

	return {
	'metas': metas,
	'inputs': inputs
	}

	except Exception as e:
	print(f"Error processing batch: {e}")
	return None


	def worker_process(questions_chunk, gpu_id, args):
	device = f'cuda:{gpu_id}'
	answer_analysis_dict = {} # Renamed from ms_answer_analysis_dict
	max_new_tokens = 256

	# --- SETUP DATALOADER ---
	batch_size = getattr(args, 'batch_size', 8)
	num_workers = getattr(args, 'num_workers', 4)

	dataset = QwenEvalDataset(questions_chunk, args)

	# Initialize the NEW single resolution collator
	collator = SingleResCollator(args.processor_path, args.patch_size)

	dataloader = DataLoader(
	dataset,
	batch_size=batch_size,
	shuffle=False,
	num_workers=num_workers,
	collate_fn=collator,
	pin_memory=True
	)

	if len(dataset) == 0: return answer_analysis_dict

	try:
	# Load Model
	print(f"Worker {gpu_id}: Loading model...")
	model_class = resolve_model_class(args.model_type)
	model = model_class.from_pretrained(
	args.model_path,
	torch_dtype=torch.bfloat16,
	trust_remote_code=True,
	device_map=device
	)
	model.eval()

	# Flags (keeping your original config)
	two_stage_flag = False
	set_twig_flags(model, args.model_type, two_stage_flag)

	# Tokenizer for decoding
	tokenizer = AutoProcessor.from_pretrained(args.processor_path, trust_remote_code=True).tokenizer

	print(f"Worker {gpu_id}: Starting inference...")

	# --- MAIN INFERENCE LOOP ---
	progress_bar = tqdm(dataloader, desc=f"GPU {gpu_id}", position=gpu_id)

	for batch_data in progress_bar:
	if batch_data is None: continue

	metas = batch_data['metas']
	inputs = batch_data['inputs'] # Now getting direct inputs, not inputs_by_res

	try:
	# Move to GPU
	inputs = inputs.to(device)

	with torch.no_grad():
	output_ids = model.generate(
	**inputs,
	max_new_tokens=max_new_tokens,
	output_attentions=False,
	pad_token_id=tokenizer.eos_token_id
	)

	# Decode
	input_ids_len = inputs.input_ids.shape[1]
	generated_ids = output_ids[:, input_ids_len:]
	batch_answers = tokenizer.batch_decode(
	generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
	)

	# --- SAVE RESULTS ---
	for idx, (meta, answer) in enumerate(zip(metas, batch_answers)):
	line_number = meta['line_number']

	tmp_sample = {
	'image': meta['image_id'],
	'prompt': meta['prompt'],
	'text': meta['ground_truth'],
	'id': meta['id'],
	'response': answer # Storing single response
	}
	answer_analysis_dict[line_number] = tmp_sample

	# Cleanup GPU memory for this batch
	del inputs, output_ids, generated_ids

	except RuntimeError as e:
	msg = "ERROR_OOM" if "out of memory" in str(e) else f"ERROR_RUNTIME: {e}"
	print(f"Worker {gpu_id}: {msg}")
	torch.cuda.empty_cache()
	# Log errors for this batch
	for meta in metas:
	line_number = meta['line_number']
	answer_analysis_dict[line_number] = {
	'id': meta['id'],
	'response': msg
	}

	except Exception as e:
	print(f"Worker {gpu_id}: Critical Error: {e}")
	import traceback
	traceback.print_exc()
	finally:
	torch.cuda.empty_cache()

	return answer_analysis_dict


	if __name__ == "__main__":
	parser = argparse.ArgumentParser(description="Parallel KL Divergence Evaluation")
	parser.add_argument('--model-path', type=str, required=True,
	help="Path to the model and processor")
	parser.add_argument('--model-type', type=str, default="auto",
	choices=["auto", "qwen2_5", "qwen2", "qwen3"],
	help="Model family for loading (auto detects from model path).")
	parser.add_argument('--processor-path', type=str, default="",
	help="Override processor path if auto detection is incorrect.")
	parser.add_argument('--question-file', type=str, default=None,
	help="Path to the JSONL question file (E.g., TextVQA)")
	parser.add_argument('--image-folder', type=str, default=None,
	help="Path to the image directory (for JSONL file)")
	parser.add_argument('--dst-save-path', type=str, required=True,
	help="Path to save the output .pkl file")
	parser.add_argument('--gpu-ids', type=str, default='0,1',
	help="Comma-separated list of GPU IDs to use (e.g., '0,1,2,3')")
	# --- ADDED: New HF Dataset args ---
	parser.add_argument('--hf_dataset', type=str, default=None,
	help="Name of the Hugging Face dataset (e.g., 'craigwu/vstar_bench')")
	parser.add_argument('--dataset_split', type=str, default='test',
	help="Split to use from the Hugging Face dataset (e.g., 'test')")
	parser.add_argument('--is_check_pred', type=bool, default=False,
	help="Whether to check high-res predictions")
	### --- END ADDITION ---
	parser.add_argument('--use-short-prompt', action='store_true',
	help="Whether to use short prompt format")
	args = parser.parse_args()
	args.model_type = infer_model_type(args.model_path, args.model_type)
	args.processor_path = resolve_processor_path(
	args.model_path, args.model_type, args.processor_path
	)
	args.patch_size = 32 if args.model_type == "qwen3" else 28


	# --- 1. Define a Wrapper Function ---
	# This helper must be defined at the top level (global scope) so it can be pickled.
	# --- 2. The Main Execution Block ---
	# Ensure start method is spawn (Crucial for PyTorch+CUDA)
	try:
	mp.set_start_method('spawn', force=True)
	except RuntimeError:
	pass

	# 1. Prepare Data and Split into Chunks
	gpu_ids = [int(gid) for gid in args.gpu_ids.split(',')]
	num_gpus = len(gpu_ids)
	questions = []

	if args.hf_dataset:
	print(f"Loading Hugging Face dataset {args.hf_dataset}, split {args.dataset_split}")
	from datasets import load_dataset

	dataset = load_dataset(args.hf_dataset, split=args.dataset_split)
	print("Casting 'image' column to load from cache...")
	questions_list = list(dataset)

	new_questions = []
	for i, tmp_ques in enumerate(questions_list):
	tmp_ques['original_line_number'] = i
	new_questions.append(tmp_ques)
	questions = new_questions
	print(f"Loaded {len(questions)} preprocessed questions from Hugging Face.")

	elif args.question_file:
	print("Loading and preprocessing questions from JSONL file...")
	if not args.image_folder:
	raise ValueError("--image-folder is required when using --question-file")

	questions_list = [json.loads(q) for q in open(os.path.expanduser(args.question_file), "r")]

	new_questions = []
	for i, tmp_ques in enumerate(questions_list):
	qs = tmp_ques['text'] if 'text' in tmp_ques else tmp_ques['prompt']
	if 'Reference OCR' in qs:
	qs = '\n'.join([ele for ele in qs.split('\n') if 'Reference OCR' not in ele])
	tmp_ques['text'] = qs
	tmp_ques['original_line_number'] = i
	new_questions.append(tmp_ques)
	questions = new_questions
	print(f"Loaded {len(questions)} preprocessed questions from local file.")

	else:
	raise ValueError("You must provide either --question-file or --hf_dataset")

	# Create chunks
	chunk_size = math.ceil(len(questions) / num_gpus)
	chunks = [questions[i:i + chunk_size] for i in range(0, len(questions), chunk_size)]

	# 2. Run processes in parallel using mp.Process (Manual Management)
	print(f"Starting parallel processing on {num_gpus} GPUs: {gpu_ids}")

	# Create a Manager to hold the results safely across processes
	manager = mp.Manager()
	return_dict = manager.dict()
	processes = []

	for i in range(num_gpus):
	chunk_data = chunks[i] if i < len(chunks) else []

	# We pass the shared 'return_dict' to the wrapper
	p = mp.Process(
	target=worker_wrapper,
	args=(chunk_data, gpu_ids[i], args, return_dict)
	)
	p.start()
	processes.append(p)

	# Wait for all processes to finish
	for p in processes:
	p.join()

	# 3. Gather and Save Results
	print("All chunks processed. Gathering results...")
	final_dict = {}

	# Extract results from the shared return_dict
	for gpu_id, sub_dict in return_dict.items():
	final_dict.update(sub_dict)

	#check save directory
	os.makedirs(os.path.dirname(os.path.expanduser(args.dst_save_path)), exist_ok=True)
	with open(os.path.expanduser(args.dst_save_path), 'wb') as f:
	pickle.dump(final_dict, f)

	print(f"✅ Successfully processed and saved {len(final_dict)} answers to {args.dst_save_path}.")


	'''
	CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
	--model-path Qwen/Qwen2.5-VL-7B-Instruct \
	--question-file /home/yuheng/datasets/textvqa/converted_llava_style_train.jsonl \
	--image-folder ../llava/playground/data/eval/textvqa/train_images \
	--dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B.pkl \
	--gpu-ids 0,1

	CUDA_VISIBLE_DEVICES=2,3 python analysis/making_base_model_response.py \
	--model-path Qwen/Qwen2.5-VL-7B-Instruct \
	--question-file /home/yuheng/datasets/vstar_coco_spatial_relation_data.jsonl \
	--image-folder /home/yuheng/datasets \
	--dst-save-path /home/yuheng/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B.pkl \
	--gpu-ids 0,1




	python analysis/making_base_model_response.py \
	--model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
	--question-file /root/epfs/sproject/datasets/textvqa/converted_llava_style_train.jsonl \
	--image-folder /root/epfs/sproject/datasets/textvqa/train_images \
	--dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/textvqa_qwen2.5vl-7B-short-prompt.pkl \
	--use-short-prompt \
	--gpu-ids 0,1

	python analysis/making_base_model_response.py \
	--model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
	--question-file /root/epfs/sproject/datasets/chartqa_28k_qa.jsonl \
	--image-folder /root/epfs/sproject/datasets \
	--dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/chartqa_qwen2.5vl-7B-short-prompt.pkl \\
	--use-short-prompt \
	--gpu-ids 0,1

	python analysis/making_base_model_response.py \
	--model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
	--question-file /root/epfs/sproject/datasets/visual_cot_infovqa_subset15k.jsonl \
	--image-folder /root/epfs/sproject/datasets \
	--dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_infovqa_qwen2.5vl-7B-short-prompt.pkl \
	--use-short-prompt \
	--gpu-ids 0,1

	python analysis/making_base_model_response.py \
	--model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
	--question-file /root/epfs/sproject/datasets/visual_cot_docvqa_subset33k.jsonl \
	--image-folder /root/epfs/sproject/datasets \
	--dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/visual_cot_docvqa_qwen2.5vl-7B-short-prompt.pkl \
	--use-short-prompt \
	--gpu-ids 0,1

	python analysis/making_base_model_response.py \
	--model-path /root/epfs/sproject/code/Qwen2.5-VL/output/qwen2_5vl-7b-roi-K18T3-185k-v1bf16-denoisedv2-twiginit-filled \
	--question-file /root/epfs/sproject/datasets/vstar_coco_spatial_relation_data.jsonl \
	--image-folder /root/epfs/sproject/datasets/ \
	--dst-save-path /root/epfs/sproject/datasets/SFT_SDRPN_PostTraining/vstar_coco_spatial_qwen2.5vl-7B-short-prompt.pkl \
	--use-short-prompt \
	--gpu-ids 0,1
	'''