scripts/hf_jobs/hf_benchmark.py

# /// script
# requires-python = ">=3.12"
# dependencies = [
#   "torch>=2.5.0",
#   "torchvision",
#   "transformers>=4.50.0",
#   "accelerate",
#   "Pillow",
#   "numpy",
# ]
# ///
"""DiffThinker HF Jobs Benchmark — Evaluate MLLM baselines on spatial reasoning.
Paper: ICML 2026, arXiv 2512.24165
Tasks: Maze 8x8 (200), Maze 32x32 (50), Sudoku 4x4 (50)
Model: Qwen2.5-VL-7B-Instruct on A10G
"""
import torch, json, os, sys, time, numpy as np
from collections import deque
from PIL import Image, ImageDraw

def bfs(grid, start, goal):
    q = deque([start]); v = {start}
    while q:
        r,c = q.popleft()
        if (r,c)==goal: return True
        for dr,dc in [(0,1),(0,-1),(1,0),(-1,0)]:
            nr,nc = r+dr,c+dc
            if 0<=nr<grid.shape[0] and 0<=nc<grid.shape[1] and grid[nr,nc]==0 and (nr,nc) not in v:
                v.add((nr,nc)); q.append((nr,nc))
    return False

def make_maze(gs, force_solvable=None):
    grid = np.zeros((gs,gs), dtype=np.uint8)
    start = (0, np.random.randint(0,gs))
    goal = (gs-1, np.random.randint(0,gs))
    for _ in range(int(gs*gs*np.random.uniform(0.12,0.30))):
        wx,wy = np.random.randint(0,gs,2)
        if (wx,wy)!=start and (wx,wy)!=goal: grid[wx,wy]=1
    has_path = bfs(grid, start, goal)
    if force_solvable is not None and has_path!=force_solvable:
        return make_maze(gs, force_solvable)
    ims = max(64, gs*4)
    img = Image.new("RGB", (ims,ims), (255,255,255))
    draw = ImageDraw.Draw(img); cw = ims//gs
    for r in range(gs):
        for c in range(gs):
            if grid[r,c]==1: draw.rectangle([c*cw,r*cw,(c+1)*cw,(r+1)*cw], fill=(100,100,100))
    draw.rectangle([start[1]*cw,start[0]*cw,(start[1]+1)*cw,(start[0]+1)*cw], fill=(0,255,0))
    draw.rectangle([goal[1]*cw,goal[0]*cw,(goal[1]+1)*cw,(goal[0]+1)*cw], fill=(255,0,0))
    return img, has_path

def make_sudoku():
    gs=4; nums=list(range(1,gs+1)); np.random.shuffle(nums)
    sol = np.zeros((gs,gs), dtype=int)
    for i in range(gs): sol[i] = np.roll(nums, i)
    puzzle = np.zeros((gs,gs), dtype=int)
    for idx in np.random.choice(gs*gs, gs*gs//2, replace=False):
        r,c = divmod(idx,gs); puzzle[r,c] = sol[r,c]
    img = Image.new("RGB", (64,64), (255,255,255))
    draw = ImageDraw.Draw(img); cw = 64//gs
    for r in range(gs):
        for c in range(gs):
            x,y = c*cw, r*cw; draw.rectangle([x,y,x+cw,y+cw], outline=(0,0,0))
            if puzzle[r,c]>0: draw.text((x+2,y+1), str(puzzle[r,c]), fill=(0,0,0))
    return img, sol, puzzle

def main():
    print("="*60)
    print("DiffThinker HF Jobs Benchmark — MLLM Baselines")
    print("Model: Qwen2.5-VL-7B-Instruct on A10G")
    print("="*60)

    device = torch.device("cuda")
    gpu = torch.cuda.get_device_name(0)
    mem = torch.cuda.get_device_properties(0).total_memory/1e9
    print(f"GPU: {gpu} ({mem:.1f}GB)")

    from transformers import AutoModelForImageTextToText, AutoProcessor
    model_id = "Qwen/Qwen2.5-VL-7B-Instruct"
    print(f"\nLoading {model_id}...")
    t0 = time.time()
    model = AutoModelForImageTextToText.from_pretrained(model_id, torch_dtype=torch.bfloat16,
                                                         device_map="auto", trust_remote_code=True)
    processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
    vr = torch.cuda.memory_allocated(0)/1e9
    print(f"Loaded in {time.time()-t0:.0f}s. VRAM: {vr:.2f}GB")

    results = {"gpu": gpu, "model": model_id, "vram_gb": round(vr,2), "tasks": []}

    for gs, ns, label in [(8, 200, "Maze 8x8"), (32, 50, "Maze 32x32")]:
        prompt = "Is there a path from green start to red goal? Answer YES or NO."
        correct = 0; lats = []
        print(f"\n--- {label} ({ns} samples) ---")
        for i in range(ns):
            img, hp = make_maze(gs)
            msgs = [{"role":"user","content":[{"type":"image","image":img},{"type":"text","text":prompt}]}]
            text = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
            inputs = processor(text=[text], images=[img], padding=True, return_tensors="pt").to(device)
            with torch.no_grad():
                torch.cuda.synchronize(); t0 = time.time()
                out = model.generate(**inputs, max_new_tokens=10, do_sample=False)
                torch.cuda.synchronize(); lats.append(time.time()-t0)
            resp = processor.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip().upper()
            correct += ("YES" in resp or "PATH" in resp) == hp
            if (i+1)%25==0: print(f"  [{i+1}/{ns}] acc={correct/(i+1)*100:.1f}% lat={np.mean(lats):.3f}s")
        acc = correct/ns*100
        print(f"  {label}: {acc:.1f}% ({correct}/{ns}), avg_lat={np.mean(lats):.3f}s")
        results["tasks"].append({"task": label, "accuracy": round(acc,1), "correct": correct,
                                  "total": ns, "avg_latency_s": round(np.mean(lats),3)})

    print("\n--- Sudoku 4x4 (50 samples) ---")
    lats = []
    for i in range(50):
        img, sol, puzzle = make_sudoku()
        msgs = [{"role":"user","content":[{"type":"image","image":img},{"type":"text","text":"Complete this 4x4 Sudoku. Output the grid."}]}]
        text = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
        inputs = processor(text=[text], images=[img], padding=True, return_tensors="pt").to(device)
        with torch.no_grad():
            torch.cuda.synchronize(); t0=time.time()
            out = model.generate(**inputs, max_new_tokens=100, do_sample=False)
            torch.cuda.synchronize(); lats.append(time.time()-t0)
        if (i+1)%10==0: print(f"  [{i+1}/50] lat={np.mean(lats):.3f}s")
    print(f"  Sudoku 4x4: avg_lat={np.mean(lats):.3f}s (accuracy requires manual review)")
    results["tasks"].append({"task":"Sudoku 4x4","avg_latency_s":round(np.mean(lats),3)})

    print("\n"+"="*60+"\nFINAL RESULTS\n"+"="*60)
    for t in results["tasks"]:
        acc = t.get("accuracy","N/A")
        print(f"  {t['task']}: acc={acc} lat={t['avg_latency_s']}s")
    print(f"\nGPU: {gpu}, VRAM: {vr:.1f}GB")

    report = {"model":model_id, "gpu":gpu, "vram_gb":round(vr,1), "results":results["tasks"]}
    print(json.dumps(report, indent=2))
    os.makedirs("/tmp/results", exist_ok=True)
    with open("/tmp/results/benchmark.json","w") as f: json.dump(report, f)

if __name__=="__main__": main()