scripts/repro_eval.py

#!/usr/bin/env python3
"""Evaluate Qwen2.5-VL-7B on DiffThinker Maze eval dataset from HF."""
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, numpy as np
from PIL import Image
from pathlib import Path

def main():
    print("=== DiffThinker: MLLM Baseline Evaluation on HF Jobs ===")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
    vram = torch.cuda.get_device_properties(0).total_memory / 1e9 if torch.cuda.is_available() else 0
    print(f"GPU: {gpu} | VRAM: {vram:.1f}GB")

    # Download eval dataset
    from huggingface_hub import snapshot_download
    print("\n[1] Downloading DiffThinker eval dataset...")
    data_dir = Path("/tmp/diffthinker_eval")
    snapshot_download("yhx12/DiffThinker_Eval", repo_type="dataset",
                      local_dir=data_dir, allow_patterns=["Maze/*"])
    maze_files = sorted((data_dir / "Maze" / "8_test").glob("*_solution.png"))
    print(f"Found {len(maze_files)} Maze 8 test samples")

    if len(maze_files) == 0:
        # Try alternative paths
        maze_files = sorted(data_dir.rglob("*solution.png"))
        print(f"Found {len(maze_files)} total solution images")

    if len(maze_files) == 0:
        # List what we have
        for p in sorted(data_dir.rglob("*"))[:30]:
            print(f"  {p.relative_to(data_dir)}")
        print("No samples found, using synthetic data")
        maze_files = []

    # Load model
    print("\n[2] Loading Qwen2.5-VL-7B-Instruct (4-bit)...")
    from transformers import (
        Qwen2_5_VLForConditionalGeneration, AutoProcessor,
        BitsAndBytesConfig
    )

    model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
    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(
        model_id, quantization_config=bnb, device_map="auto",
        torch_dtype=torch.bfloat16, trust_remote_code=True,
    )
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    print(f"Model loaded | B parameters")

    # Eval on Maze 8x8
    print("\n[3] Evaluating on Maze 8x8 samples...")
    prompt = "Solve this maze. Green=start, red=goal, gray=walls. Output the path as list of (row,col) coordinates."

    results = []
    for i, sol_path in enumerate(maze_files[:10]):
        task_img = sol_path.parent / sol_path.name.replace("_solution", "")
        if not task_img.exists():
            task_img = sol_path
        img = Image.open(task_img)
        print(f"  Sample {i+1}: {task_img.name} ({img.size})")

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

        torch.cuda.synchronize()
        t0 = time.time()
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=256, do_sample=False, temperature=1.0)
        torch.cuda.synchronize()
        lat = time.time() - t0
        resp = processor.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        print(f"    Latency: {lat:.2f}s | Resp: {resp[:120]}...")
        results.append({"sample": task_img.name, "latency": round(lat, 2), "response": resp[:200]})

    # Analysis
    print("\n[4] Summary")
    avg_lat = sum(r["latency"] for r in results) / len(results) if results else 0
    print(json.dumps({
        "model": "Qwen2.5-VL-7B-Instruct (4-bit)",
        "gpu": gpu, "vram_gb": round(vram, 1),
        "samples": len(results),
        "avg_latency_s": round(avg_lat, 2),
        "results": results,
    }, indent=2))
    return 0

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