| |
| """Unified vLLM evaluator for image-sequence multiple-choice benchmarks. |
| |
| Supported benchmarks: |
| - MMSI-Bench: TSV/parquet with base64 image list in `image`. |
| - MindCube-Tiny: official 1,050-sample Hugging Face split. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import ast |
| import base64 |
| import csv |
| import io |
| import json |
| import os |
| import re |
| import sys |
| import time |
| from collections import defaultdict |
| from typing import Any, Dict, List |
|
|
| from canonical_data import load_json_records |
| from mindcube.data_utils import ( |
| OFFICIAL_DATA_SOURCE, |
| decode_mindcube_images, |
| load_mindcube_records, |
| mindcube_answer_value, |
| mindcube_group, |
| mindcube_prompt, |
| ) |
| from PIL import Image |
|
|
| CHOICES = list("ABCDEFGH") |
| ANSWER_RE = re.compile(r"<answer>\s*(.*?)\s*</answer>", re.DOTALL | re.IGNORECASE) |
| TAG_RE = re.compile(r"</?think>", re.IGNORECASE) |
| FORMAT_PRIORITY = { |
| "start": 10, |
| "end": 9, |
| "phrase": 7, |
| "parentheses": 6, |
| "period": 5, |
| "colon": 4, |
| "right_paren": 3, |
| "space": 2, |
| "fallback": 0, |
| } |
| ANSWER_PHRASES = [ |
| "the answer is", |
| "answer is", |
| "the correct answer is", |
| "correct answer is", |
| "the best answer is", |
| "best answer is", |
| "the correct option is", |
| "correct option is", |
| "i choose", |
| "i select", |
| "my answer is", |
| "答案是", |
| "答案为", |
| ] |
| PROMPT_TAIL = ( |
| "Choose the best answer from the options. " |
| "Put exactly one uppercase option letter inside <answer>...</answer> " |
| "Do not explain. Example: <answer>A</answer>" |
| ) |
|
|
|
|
| def strip_think_block(text: str) -> str: |
| text = (text or "").strip() |
| if not text: |
| return "" |
| parts = re.split(r"</think>", text, flags=re.IGNORECASE) |
| if len(parts) > 1: |
| text = parts[-1] |
| else: |
| text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE) |
| return TAG_RE.sub("", text).strip() |
|
|
|
|
| def strip_answer_tags(text: str) -> str: |
| matches = ANSWER_RE.findall(text or "") |
| if matches: |
| return matches[-1].strip() |
| return (text or "").strip() |
|
|
|
|
| def extract_mcq_answer(response: str, choices: List[str] | None = None) -> str: |
| if not response or not response.strip(): |
| return "" |
| all_choices = choices or CHOICES |
| text = strip_answer_tags(strip_think_block(response)).strip() |
| if not text: |
| return "" |
| for char in [",", ".", "!", "?", ";", ":", "'", '"', "。", "："]: |
| text = text.strip(char) |
| padded = " " + text + " " |
| candidates = [] |
|
|
| for ch in all_choices: |
| for token, fmt in ( |
| (f"({ch})", "parentheses"), |
| (f"{ch}.", "period"), |
| (f"{ch}:", "colon"), |
| (f"{ch})", "right_paren"), |
| (f"{ch} ", "space"), |
| ): |
| pos = padded.rfind(token) |
| if pos != -1: |
| candidates.append((ch, pos, fmt)) |
|
|
| lower = padded.lower() |
| for phrase in ANSWER_PHRASES: |
| idx = lower.rfind(phrase.lower()) |
| if idx != -1: |
| after = idx + len(phrase) |
| for ch in all_choices: |
| m = re.search(rf"\b{re.escape(ch)}\b", padded[after:], flags=re.IGNORECASE) |
| if m: |
| candidates.append((ch, after + m.start(), "phrase")) |
|
|
| stripped = padded.strip() |
| for ch in all_choices: |
| if stripped.upper() == ch: |
| candidates.append((ch, 0, "start")) |
| elif stripped.startswith(ch) and (len(stripped) == 1 or not stripped[1].isalpha()): |
| candidates.append((ch, 0, "start")) |
| elif stripped.endswith(ch) and (len(stripped) == 1 or not stripped[-2].isalpha()): |
| candidates.append((ch, len(padded) - 1, "end")) |
|
|
| if not candidates: |
| for ch in all_choices: |
| m = re.search(rf"\b{re.escape(ch)}\b", padded) |
| if m: |
| candidates.append((ch, m.start(), "fallback")) |
| if not candidates: |
| return "" |
| candidates.sort(key=lambda x: (FORMAT_PRIORITY.get(x[2], 0), x[1]), reverse=True) |
| return candidates[0][0] |
|
|
|
|
| def load_records(args) -> List[Dict[str, Any]]: |
| if args.bench == "mindcube": |
| return load_mindcube_records( |
| args.data_file, |
| chunk=args.chunk, |
| index=args.index, |
| expected_samples=args.expected_samples, |
| cache_dir=args.hf_cache_dir or None, |
| ) |
| if args.data_file.endswith((".jsonl", ".json")): |
| return load_json_records(args.data_file) |
| if args.data_file.endswith(".parquet"): |
| import pandas as pd |
|
|
| return pd.read_parquet(args.data_file).to_dict("records") |
|
|
| csv.field_size_limit(sys.maxsize) |
| with open(args.data_file, encoding="utf-8", errors="replace", newline="") as f: |
| return list(csv.DictReader(f, delimiter="\t")) |
|
|
|
|
| def shard(records: List[Dict[str, Any]], chunk: int, index: int) -> List[Dict[str, Any]]: |
| if chunk <= 1: |
| return records |
| return records[index::chunk] |
|
|
|
|
| def decode_mmsi_images(value: Any, max_images: int) -> List[Image.Image]: |
| images_b64 = ast.literal_eval(value) if isinstance(value, str) else value |
| if not isinstance(images_b64, (list, tuple)): |
| images_b64 = [images_b64] |
| images = [] |
| for item in images_b64[:max_images]: |
| if not item: |
| continue |
| path = "" |
| if isinstance(item, dict): |
| path = str( |
| item.get("path") |
| or item.get("image") |
| or item.get("image_path") |
| or "" |
| ) |
| elif isinstance(item, str) and os.path.isfile(item): |
| path = item |
| if path: |
| images.append(Image.open(path).convert("RGB")) |
| continue |
| raw = item if isinstance(item, bytes) else base64.b64decode(str(item)) |
| images.append(Image.open(io.BytesIO(raw)).convert("RGB")) |
| return images |
|
|
|
|
| def normalize_choices(value: Any) -> List[str]: |
| if value is None: |
| return [] |
| if hasattr(value, "tolist"): |
| value = value.tolist() |
| if isinstance(value, (list, tuple)): |
| return [str(v) for v in value] |
| return [str(value)] |
|
|
|
|
| def answer_choices(choices: List[Any] | None) -> List[str]: |
| n = len(choices) if choices is not None else 0 |
| return CHOICES[:n] if n > 0 else list("ABCD") |
|
|
|
|
| def choices_from_question(question: str) -> List[str]: |
| found = re.findall(r"(?:^|[\s,])([A-H])\s*[:\).]", question or "") |
| out = [] |
| for ch in found: |
| if ch not in out: |
| out.append(ch) |
| return out or list("ABCD") |
|
|
|
|
| def build_prompt(args, rec: Dict[str, Any]) -> tuple[str, List[str]]: |
| if args.bench == "mmsi": |
| question = str(rec.get("question") or "").strip() |
| labels = choices_from_question(question) |
| return f"{question}\n{PROMPT_TAIL}", labels |
|
|
| return mindcube_prompt(rec, PROMPT_TAIL) |
|
|
|
|
| def get_answer(args, rec: Dict[str, Any]) -> str: |
| answer = ( |
| mindcube_answer_value(rec) |
| if args.bench == "mindcube" |
| else rec.get("answer") |
| ) |
| if answer is None: |
| return "" |
| if args.bench == "mindcube" and isinstance(answer, (int, float)) and not isinstance(answer, bool): |
| idx = int(answer) |
| labels = answer_choices(normalize_choices(rec.get("choices"))) |
| if 0 <= idx < len(labels): |
| return labels[idx] |
| return extract_mcq_answer(str(answer), CHOICES) or str(answer).strip().upper()[:1] |
|
|
|
|
| def get_group_key(args, rec: Dict[str, Any]) -> str: |
| if args.bench == "mmsi": |
| return str(rec.get("category") or "unknown") |
| return mindcube_group(rec) |
|
|
|
|
| def get_filter_key(args) -> str: |
| return "category" if args.bench == "mmsi" else "task" |
|
|
|
|
| def render_chat_prompt(messages, processor, enable_thinking: bool = False) -> str: |
| return processor.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=enable_thinking, |
| ) |
|
|
|
|
| def prepare_for_vllm(images: List[Image.Image], prompt: str, processor, args): |
| from qwen_vl_utils import process_vision_info |
|
|
| content: List[Dict[str, Any]] = [] |
| for image in images: |
| item: Dict[str, Any] = { |
| "type": "image", |
| "image": image, |
| "max_pixels": args.image_max_pixels, |
| } |
| if args.image_min_pixels > 0: |
| item["min_pixels"] = args.image_min_pixels |
| content.append(item) |
| content.append({"type": "text", "text": prompt}) |
| messages = [{"role": "user", "content": content}] |
| text = render_chat_prompt(messages, processor, args.enable_thinking) |
| image_inputs, _video_inputs = process_vision_info(messages, image_patch_size=args.patch_size) |
| return { |
| "prompt": text, |
| "multi_modal_data": {"image": image_inputs}, |
| "mm_processor_kwargs": {"do_resize": False}, |
| } |
|
|
|
|
| def pct(correct: int, total: int) -> float: |
| return round(100.0 * correct / total, 2) if total else 0.0 |
|
|
|
|
| def aggregate(args, results: List[Dict[str, Any]]) -> Dict[str, Any]: |
| correct = sum(int(r.get("score", 0)) for r in results) |
| parsed = sum(1 for r in results if r.get("pred_answer")) |
| groups = defaultdict(lambda: {"correct": 0, "total": 0}) |
| for r in results: |
| groups[str(r.get("group") or "unknown")]["total"] += 1 |
| groups[str(r.get("group") or "unknown")]["correct"] += int(r.get("score", 0)) |
| group_name = "by_category" if args.bench == "mmsi" else "by_task" |
| return { |
| "num_samples": len(results), |
| "correct": correct, |
| "accuracy": pct(correct, len(results)), |
| "parse_rate": pct(parsed, len(results)), |
| group_name: { |
| k: {"accuracy": pct(v["correct"], v["total"]), "correct": v["correct"], "total": v["total"]} |
| for k, v in sorted(groups.items()) |
| }, |
| } |
|
|
|
|
| def decode_images_for_record(args, rec: Dict[str, Any]) -> List[Image.Image]: |
| if args.bench == "mmsi": |
| return decode_mmsi_images( |
| rec.get("images") or rec.get("image"), |
| args.max_images, |
| ) |
| return decode_mindcube_images(rec, args.max_images, args.min_image_bytes) |
|
|
|
|
| def evaluate(llm, sampling_params, processor, args) -> Dict[str, Any]: |
| records = load_records(args) |
| filter_value = args.category if args.bench == "mmsi" else args.task |
| if filter_value: |
| if args.bench == "mindcube": |
| records = [ |
| record |
| for record in records |
| if mindcube_group(record) == filter_value |
| ] |
| else: |
| key = get_filter_key(args) |
| records = [ |
| record |
| for record in records |
| if str(record.get(key) or "") == filter_value |
| ] |
| if args.max_samples and args.max_samples > 0: |
| records = records[: args.max_samples] |
| if args.bench == "mmsi": |
| records = shard(records, args.chunk, args.index) |
| print(f"Loaded {len(records)} {args.bench} samples for shard {args.index}/{args.chunk}", flush=True) |
|
|
| results: List[Dict[str, Any]] = [] |
| t0 = time.time() |
| for start in range(0, len(records), args.batch_size): |
| batch = records[start:start + args.batch_size] |
| inputs = [] |
| keep = [] |
| prompts: Dict[int, str] = {} |
| labels_by_j: Dict[int, List[str]] = {} |
| for j, rec in enumerate(batch): |
| try: |
| images = decode_images_for_record(args, rec) |
| if not images: |
| raise ValueError("no decoded images") |
| prompt, labels = build_prompt(args, rec) |
| inputs.append(prepare_for_vllm(images, prompt, processor, args)) |
| keep.append(j) |
| prompts[j] = prompt |
| labels_by_j[j] = labels |
| except Exception as e: |
| print(f"[warn] skip sample {start + j}: {type(e).__name__}: {e}", flush=True) |
|
|
| outputs = llm.generate(inputs, sampling_params=sampling_params) if inputs else [] |
| out_by_j = {j: out for j, out in zip(keep, outputs)} |
|
|
| for j, rec in enumerate(batch): |
| if j not in out_by_j: |
| continue |
| raw = out_by_j[j].outputs[0].text |
| labels = labels_by_j[j] |
| pred = extract_mcq_answer(raw, labels) |
| gt = get_answer(args, rec) |
| score = 1.0 if pred and gt and pred.upper() == gt.upper() else 0.0 |
| results.append({ |
| "id": str(rec.get("index") or rec.get("id") or rec.get("sample_id") or start + j), |
| "bench": args.bench, |
| "group": get_group_key(args, rec), |
| "question": rec.get("question"), |
| "prompt": prompts.get(j, ""), |
| "answer": gt, |
| "pred_answer": pred, |
| "raw_prediction": raw, |
| "score": score, |
| }) |
|
|
| done = min(start + args.batch_size, len(records)) |
| if done % 200 == 0 or done == len(records): |
| elapsed = max(time.time() - t0, 1e-6) |
| summary = aggregate(args, results) |
| print( |
| f"[{done}/{len(records)}] {elapsed:.1f}s " |
| f"acc={summary['accuracy']:.2f}% parse={summary['parse_rate']:.2f}%", |
| flush=True, |
| ) |
|
|
| full_mindcube_eval = ( |
| args.bench == "mindcube" |
| and args.chunk == 1 |
| and not args.task |
| and not args.max_samples |
| ) |
| if ( |
| full_mindcube_eval |
| and args.expected_samples > 0 |
| and len(results) != args.expected_samples |
| ): |
| raise RuntimeError( |
| f"Expected {args.expected_samples} official MindCube-Tiny " |
| f"results, got {len(results)}." |
| ) |
|
|
| out_name = f"results_{args.bench}" |
| if args.chunk > 1: |
| out_name += f"_shard{args.index}" |
| out_path = os.path.join(args.output_dir, out_name + ".json") |
| with open(out_path, "w", encoding="utf-8") as f: |
| json.dump(results, f, ensure_ascii=False, indent=2) |
|
|
| summary = aggregate(args, results) |
| if args.bench == "mindcube": |
| summary["data_source"] = args.data_file |
| if full_mindcube_eval and args.expected_samples > 0: |
| summary["expected_samples"] = args.expected_samples |
| summary["coverage"] = round( |
| 100.0 * len(results) / args.expected_samples, 2 |
| ) |
| summary_name = f"summary_shard{args.index}.json" if args.chunk > 1 else "summary.json" |
| with open(os.path.join(args.output_dir, summary_name), "w", encoding="utf-8") as f: |
| json.dump(summary, f, ensure_ascii=False, indent=2) |
| print(json.dumps(summary, ensure_ascii=False, indent=2), flush=True) |
| return summary |
|
|
|
|
| def parse_args(): |
| p = argparse.ArgumentParser(description="Unified image-sequence MC evaluation via vLLM.") |
| p.add_argument("--bench", required=True, choices=["mmsi", "mindcube"]) |
| p.add_argument("--model_path", required=True) |
| p.add_argument("--processor_path", default="") |
| p.add_argument("--data_file", default="") |
| p.add_argument("--output_dir", required=True) |
| p.add_argument("--image_min_pixels", type=int, default=4096) |
| p.add_argument("--image_max_pixels", type=int, default=262144) |
| p.add_argument("--max_images", type=int, default=8) |
| p.add_argument("--min_image_bytes", type=int, default=1024) |
| p.add_argument("--patch_size", type=int, default=16) |
| p.add_argument("--tensor_parallel_size", type=int, default=1) |
| p.add_argument("--gpu_memory_utilization", type=float, default=0.95) |
| p.add_argument("--max_model_len", type=int, default=65536) |
| p.add_argument("--max_new_tokens", type=int, default=64) |
| p.add_argument("--max_num_batched_tokens", type=int, default=65536) |
| p.add_argument("--batch_size", type=int, default=16) |
| p.add_argument("--max_samples", type=int, default=0) |
| p.add_argument("--expected_samples", type=int, default=1050) |
| p.add_argument("--hf_cache_dir", default="") |
| p.add_argument("--category", default="") |
| p.add_argument("--task", default="") |
| p.add_argument("--temperature", type=float, default=0.0) |
| p.add_argument("--top_p", type=float, default=1.0) |
| p.add_argument("--top_k", type=int, default=-1) |
| p.add_argument("--chunk", type=int, default=1) |
| p.add_argument("--index", type=int, default=0) |
| p.add_argument("--enable_thinking", action="store_true") |
| return p.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| if args.bench == "mindcube" and not args.data_file: |
| args.data_file = OFFICIAL_DATA_SOURCE |
| if not args.data_file: |
| raise ValueError("--data_file is required for MMSI-Bench") |
| args.processor_path = args.processor_path or args.model_path |
| os.makedirs(args.output_dir, exist_ok=True) |
|
|
| from transformers import AutoProcessor, AutoTokenizer |
| from vllm import LLM, SamplingParams |
|
|
| processor = AutoProcessor.from_pretrained( |
| args.processor_path, |
| padding_side="left", |
| trust_remote_code=True, |
| min_pixels=args.image_min_pixels, |
| max_pixels=args.image_max_pixels, |
| ) |
| tokenizer = AutoTokenizer.from_pretrained(args.processor_path, trust_remote_code=True) |
| tokenizer.padding_side = "left" |
| processor.tokenizer = tokenizer |
|
|
| llm = LLM( |
| model=args.model_path, |
| tensor_parallel_size=args.tensor_parallel_size, |
| gpu_memory_utilization=args.gpu_memory_utilization, |
| max_model_len=args.max_model_len, |
| max_num_batched_tokens=args.max_num_batched_tokens, |
| trust_remote_code=True, |
| limit_mm_per_prompt={"image": args.max_images}, |
| ) |
| sampling = SamplingParams( |
| max_tokens=args.max_new_tokens, |
| temperature=args.temperature, |
| top_p=args.top_p, |
| top_k=args.top_k, |
| ) |
| evaluate(llm, sampling, processor, args) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
