scripts/repro_eval2.py

#!/usr/bin/env python3
"""Proper accuracy eval: Qwen2.5-VL-7B on DiffThinker Maze 8x8 with ground-truth comparison."""
import subprocess, sys
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q",
    "torch>=2.1.0", "torchvision", "transformers>=4.49.0", "accelerate",
    "bitsandbytes", "Pillow", "numpy", "safetensors", "huggingface_hub"])

import json, os, time, torch, re, numpy as np
from PIL import Image
from pathlib import Path
from huggingface_hub import snapshot_download

def parse_path(resp):
    """Extract (row,col) coordinates from model response."""
    coords = re.findall(r'[\(\[][\s]*(\d+)\s*[,\s]\s*(\d+)[\)\]]', resp)
    if coords:
        return [(int(r), int(c)) for r, c in coords]
    return []

def load_solution(sol_path):
    """Extract path from solution image (red=path, green=start, blue=path)."""
    img = Image.open(sol_path).convert("RGB")
    arr = np.array(img)
    grid_size = 8
    cells = arr.shape[0] // grid_size
    path = []
    for r in range(grid_size):
        for c in range(grid_size):
            y, x = r * cells + cells//2, c * cells + cells//2
            px = arr[y, x]
            # Red or blue path pixels
            if px[0] > 200 and px[1] < 100:
                path.append((r, c))
            elif px[0] < 100 and px[1] < 100 and px[2] > 200:
                path.append((r, c))
    return path

def main():
    print("=== DiffThinker: MLLM Maze Accuracy Evaluation ===")
    import torch
    device = torch.device("cuda")
    gpu = torch.cuda.get_device_name(0)
    vram = torch.cuda.get_device_properties(0).total_memory / 1e9

    # Download eval data
    print("\n[1] Loading DiffThinker eval dataset (Maze 8x8)...")
    data_dir = Path("/tmp/dt_eval")
    snapshot_download("yhx12/DiffThinker_Eval", repo_type="dataset",
                      local_dir=data_dir, allow_patterns=["Maze/8_test/*"])
    files = sorted((data_dir / "Maze" / "8_test").glob("*_solution.png"))
    print(f"Found {len(files)} Maze 8x8 eval samples")

    # Load model
    print("\n[2] Loading Qwen2.5-VL-7B-Instruct (4-bit)...")
    from transformers import (
        Qwen2_5_VLForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
    )
    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                              bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True)
    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
        "Qwen/Qwen2.5-VL-7B-Instruct", quantization_config=bnb, device_map="cuda",
        torch_dtype=torch.bfloat16, trust_remote_code=True)
    proc = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct", trust_remote_code=True)

    prompt = ("Solve this 8x8 maze. Green=start, red=goal, gray=walls. "
              "Output ONLY the path as a list of (row,column) coordinates from start to goal.")

    # Eval on all samples
    print(f"\n[3] Evaluating on {len(files)} samples...")
    results = []
    for sol_path in files[:20]:
        img_path = Path(str(sol_path).replace("_solution.png", ".png"))
        if not img_path.exists():
            img_path = sol_path
        img = Image.open(img_path)

        # Get expected path from solution image
        expected = load_solution(sol_path)

        msg = [{"role": "user", "content": [
            {"type": "image", "image": img}, {"type": "text", "text": prompt}
        ]}]
        text = proc.apply_chat_template(msg, tokenize=False, add_generation_prompt=True)
        inputs = proc(text=[text], images=[img], padding=True, return_tensors="pt").to(device)

        t0 = time.time()
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=256, do_sample=False)
        lat = time.time() - t0
        resp = proc.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)

        # Parse path and check accuracy
        pred_path = parse_path(resp)
        # Check if path connects start to goal
        correct = 0
        if pred_path and expected:
            overlap = len(set(pred_path) & set(expected))
            correct = overlap / max(len(expected), 1)

        results.append({
            "sample": img_path.name, "latency": round(lat, 2),
            "pred_len": len(pred_path), "expected_len": len(expected),
            "overlap": f"{overlap}/{len(expected)}" if expected else "N/A",
            "overlap_pct": round(correct * 100, 1),
        })
        print(f"  {img_path.name}: {overlap}/{len(expected)} path overlap ({correct*100:.0f}%), {lat:.1f}s")

    # Summary
    print("\n[4] RESULTS")
    correct_runs = sum(1 for r in results if r["expected_len"] > 0 and r["overlap_pct"] >= 50)
    total_runs = sum(1 for r in results if r["expected_len"] > 0)
    avg_lat = sum(r["latency"] for r in results) / len(results)
    out = {
        "model": "Qwen2.5-VL-7B-Instruct (4-bit)",
        "gpu": gpu, "vram_gb": round(vram, 1),
        "samples": len(results),
        "samples_with_overlap_ge50pct": correct_runs,
        "total_comparable": total_runs,
        "accuracy_pct": round(correct_runs / total_runs * 100, 1) if total_runs else 0,
        "avg_latency_s": round(avg_lat, 2),
    }
    print(json.dumps(out, indent=2))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())