scripts/vllm_vl_maze.py

#!/usr/bin/env python3
"""Qwen2.5-VL-7B on Modal A10 (24GB) — full bf16, no bitsandbytes.
Evaluates vision-centric maze understanding for Claims 2 & 3.
Paper: 2512.24165
"""
import modal, json, os, time, io
import numpy as np
from PIL import Image, ImageDraw

app = modal.App("diffthinker-vl-maze")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("torch>=2.1.0", "torchvision", "transformers>=4.49.0",
                 "accelerate", "Pillow", "numpy", "huggingface_hub")
)

def make_maze(grid_size=8, force_solvable=False):
    g = np.zeros((grid_size, grid_size), dtype=int)
    attempts = 0
    while True:
        g.fill(0)
        start = (0, np.random.randint(0, grid_size))
        goal = (grid_size-1, np.random.randint(0, grid_size))
        for _ in range(int(grid_size*grid_size*0.12)):
            wx, wy = np.random.randint(1, grid_size-1, 2)
            g[wx, wy] = 1
        from collections import deque
        q = deque([start])
        parent = {start: None}
        solved = False
        while q:
            r, c = q.popleft()
            if (r, c) == goal:
                solved = True; break
            for dr, dc in [(0,1),(0,-1),(1,0),(-1,0)]:
                nr, nc = r+dr, c+dc
                if 0 <= nr < grid_size and 0 <= nc < grid_size and g[nr,nc] != 1 and (nr,nc) not in parent:
                    parent[(nr,nc)] = (r,c)
        if force_solvable and not solved:
            attempts += 1
            if attempts > 50:
                break
            continue
        break
    img = Image.new("RGB", (grid_size*40, grid_size*40), "white")
    draw = ImageDraw.Draw(img)
    for r in range(grid_size):
        for c in range(grid_size):
            x, y = c*40, r*40
            if g[r,c] == 1:
                draw.rectangle([x, y, x+40, y+40], fill="gray")
            draw.rectangle([x, y, x+40, y+40], outline="black", width=1)
    sx, sy = start[1]*40+20, start[0]*40+20
    gx, gy = goal[1]*40+20, goal[0]*40+20
    draw.ellipse([sx-8, sy-8, sx+8, sy+8], fill="green")
    draw.ellipse([gx-8, gy-8, gx+8, gy+8], fill="red")
    from collections import deque
    q = deque([start])
    parent = {start: None}
    solved = False
    while q:
        r, c = q.popleft()
        if (r,c) == goal:
            solved = True; break
        for dr, dc in [(0,1),(0,-1),(1,0),(-1,0)]:
            nr, nc = r+dr, c+dc
            if 0 <= nr < grid_size and 0 <= nc < grid_size and g[nr,nc] != 1 and (nr,nc) not in parent:
                parent[(nr,nc)] = (r,c)
    return img, solved

vol = modal.Volume.from_name("hf-hub-cache", create_if_missing=True)

@app.cls(image=image, gpu="A10G", secrets=[modal.Secret.from_name("huggingface-secret")],
         volumes={"/root/.cache/huggingface": vol}, scaledown_window=60, timeout=900)
class VLMBaseline:
    @modal.enter()
    def load(self):
        from transformers import AutoModelForImageTextToText, AutoProcessor
        import torch
        self.model = AutoModelForImageTextToText.from_pretrained(
            "Qwen/Qwen2.5-VL-7B-Instruct",
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self.processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
        print(f"Loaded. VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

    @modal.method()
    def eval_maze_vision(self, num_test=12):
        import torch, random
        # Build balanced set: half solvable, half unsolvable
        test_items = []
        for _ in range(num_test // 2):
            test_items.append(make_maze(force_solvable=True))
        for _ in range(num_test - num_test // 2):
            test_items.append(make_maze(force_solvable=False))
        random.shuffle(test_items)
        correct = 0
        out_results = []
        for i, (img, solved) in enumerate(test_items):
            messages = [
                {"role": "user", "content": [
                    {"type": "image", "image": img},
                    {"type": "text", "text": "Does a path exist from green start to red goal in this maze? Answer YES or NO."},
                ]}
            ]
            text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = self.processor(text=[text], images=[img], padding=True, return_tensors="pt").to("cuda")
            t0 = time.time()
            with torch.no_grad():
                out = self.model.generate(**inputs, max_new_tokens=64, do_sample=False)
            lat = time.time() - t0
            response = self.processor.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
            is_correct = (("yes" in response.lower()) == solved)
            correct += 1 if is_correct else 0
            out_results.append({"idx": i, "latency": round(lat, 3), "correct": is_correct, "solved": solved, "response": response[:100]})
            print(f"  [{i+1}] lat={lat:.3f}s correct={is_correct} (solved={solved})")
        acc = correct / num_test * 100
        avg_lat = sum(r["latency"] for r in out_results) / num_test
        print(f"\nVL vision acc={acc:.1f}% ({correct}/{num_test}) lat={avg_lat:.3f}s")
        return {"accuracy": acc, "correct": correct, "total": num_test, "avg_latency": round(avg_lat, 3), "results": out_results}

@app.local_entrypoint()
def main():
    vlm = VLMBaseline()
    res = vlm.eval_maze_vision.remote(num_test=20)
    print(json.dumps(res, indent=2))