fpqa_modal.py

"""FloorplanQA eval on Modal with vLLM server."""
import json, os, time, subprocess, sys, copy, random
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

import modal

app = modal.App("fpqa-eval")

vllm_image = (
    modal.Image.from_registry("nvidia/cuda:12.9.0-devel-ubuntu22.04", add_python="3.12")
    .entrypoint([])
    .uv_pip_install("vllm==0.21.0", "shapely", "huggingface_hub", "numpy", "openai")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

MODEL = "Qwen/Qwen2.5-7B-Instruct"
QTYPES = ["pair_distance","view_angle","free_space","obstruction","placement","repositioning"]
CATS = {"pair_distance":"Metric","view_angle":"Metric","free_space":"Topology",
        "placement":"Topology","obstruction":"Topology","repositioning":"Dynamic"}
SAMPLE = 5


@app.cls(gpu="A10G:1", image=vllm_image, timeout=7200, scaledown_window=600)
class FPQAEval:
    model_name: str = modal.parameter(default="Qwen/Qwen2.5-7B-Instruct")

    @modal.enter()
    def start_vllm(self):
        os.environ["VLLM_USE_V1"] = "1"
        cmd = [
            sys.executable, "-m", "vllm.entrypoints.openai.api_server",
            "--model", self.model_name,
            "--port", "8000",
            "--max-model-len", "8192",
            "--enforce-eager",
        ]
        self.proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        import urllib.request
        for _ in range(300):
            try:
                urllib.request.urlopen("http://localhost:8000/v1/models", timeout=2)
                print(f"vLLM ready with {self.model_name}", flush=True)
                return
            except:
                time.sleep(3)
        raise RuntimeError(f"vLLM failed to start for {self.model_name}")

    @modal.exit()
    def stop_vllm(self):
        if self.proc:
            self.proc.terminate()
            self.proc.wait(30)

    def _query(self, messages):
        from openai import OpenAI
        c = OpenAI(base_url="http://localhost:8000/v1", api_key="x")
        r = c.chat.completions.create(
            model=self.model_name,
            messages=messages,
            max_tokens=2000,
            temperature=0.0,
        )
        return r.choices[0].message.content

    @modal.method()
    def evaluate(self, sample=SAMPLE, fmt="json", perturb=0.0, out_dir="/results"):
        from huggingface_hub import snapshot_download
        from shapely.geometry import Polygon, LineString
        from shapely.ops import unary_union
        import shapely.affinity

        layout_root = snapshot_download(
            "OldDelorean/FloorplanQA-Layouts", repo_type="dataset",
            ignore_patterns=["*.md", "*.txt"]
        )
        layout_dir = os.path.join(layout_root, "layouts")
        print(f"Layouts downloaded to {layout_dir}", flush=True)

        def pts(c):
            return [(p["x"], p["y"]) for p in c] if c else []

        def poly(p):
            if len(p) < 3:
                return None
            try:
                P = Polygon(p)
                return P if P.is_valid else None
            except:
                return None

        def room_poly(d):
            return poly(pts(d.get("room_boundary", [])))

        def obs_polys(d):
            r = []
            for o in d.get("objects", []):
                lab = (o.get("label", "") or "").lower()
                if any(x in lab for x in ["rug", "light", "ceiling", "vent"]):
                    continue
                p = poly(pts(o.get("points", [])))
                if p:
                    r.append(p)
            return r

        def gen_gt(qt, data):
            room = room_poly(data)
            if not room or room.is_empty:
                return None
            obs = obs_polys(data)

            if qt == "pair_distance":
                objs = [(o.get("label", ""), poly(pts(o.get("points", [])))) for o in data.get("objects", [])]
                objs = [(l, p) for l, p in objs if p]
                if len(objs) < 2:
                    return None
                d = round(objs[0][1].centroid.distance(objs[1][1].centroid), 4)
                return {"a": objs[0][0], "b": objs[1][0], "answer": d}
            elif qt == "view_angle":
                objs = [(o.get("label", ""), poly(pts(o.get("points", [])))) for o in data.get("objects", [])]
                objs = [(l, p) for l, p in objs if p]
                if len(objs) < 2:
                    return None
                import math
                ax, ay = objs[0][1].centroid.x, objs[0][1].centroid.y
                bx, by = objs[1][1].centroid.x, objs[1][1].centroid.y
                angle = round(math.degrees(math.atan2(by - ay, bx - ax)), 1)
                if angle < 0:
                    angle += 360
                return {"a": objs[0][0], "b": objs[1][0], "answer": angle}
            elif qt == "free_space":
                free = room.difference(unary_union(obs)) if obs else room
                pct = round(free.area / room.area * 100, 1) if room.area > 0 else 0
                return {"answer": pct}
            elif qt == "obstruction":
                objs = [(o.get("label", ""), poly(pts(o.get("points", [])))) for o in data.get("objects", [])]
                objs = [(l, p) for l, p in objs if p]
                if len(objs) < 2:
                    return None
                line = LineString([objs[0][1].centroid.coords[0], objs[1][1].centroid.coords[0]])
                blocked = [objs[i][0] for i in range(2, len(objs)) if line.intersects(objs[i][1])]
                return {"a": objs[0][0], "b": objs[1][0], "answer": blocked if blocked else ["none"]}
            elif qt == "placement":
                objs = [(o.get("label", ""), poly(pts(o.get("points", [])))) for o in data.get("objects", [])]
                objs = [(l, p) for l, p in objs if p]
                if not objs:
                    return None
                wz = room.boundary.buffer(0.5)
                return {"answer": {lab: ("wall" if wz.contains(p.centroid) else "center") for lab, p in objs}}
            elif qt == "repositioning":
                objs = [(o.get("label", ""), poly(pts(o.get("points", [])))) for o in data.get("objects", [])]
                objs = [(l, p) for l, p in objs if p]
                if not objs:
                    return None
                moved = shapely.affinity.translate(objs[0][1], xoff=1.0)
                others = unary_union([objs[i][1] for i in range(1, len(objs))])
                collides = moved.intersects(others) if others else False
                return {"a": objs[0][0], "answer": "collision" if collides else "no collision"}

        def make_prompt(qt, gt, data, fmt):
            if fmt == "json":
                s = json.dumps({k: v for k, v in data.items() if k != "objects"}, indent=1, default=str)[:2000]
                header = f"Floor Layout (JSON):\n{s}"
            else:
                s = "<layout>\n"
                for k, v in data.items():
                    if k == "objects":
                        continue
                    s += f"  <{k}>{v}</{k}>\n"
                s += "</layout>"
                header = f"Floor Layout (XML):\n{s}"
            qmap = {
                "pair_distance": f"What is the distance between {gt.get('a','object A')} and {gt.get('b','object B')}?",
                "view_angle": f"What is the angle from {gt.get('a','object A')} to {gt.get('b','object B')}?",
                "free_space": "What percentage of the room is free space?",
                "obstruction": f"What objects obstruct the path from {gt.get('a','object A')} to {gt.get('b','object B')}?",
                "placement": "For each object, is it near the wall or in the center?",
                "repositioning": f"If you move {gt.get('a','an object')} 1 meter to the right, does it collide with any other object?"
            }
            return f"{header}\n\nQuestion: {qmap.get(qt,'?')}\n\nReply with just the answer."

        def score(qt, gt, txt):
            import re
            if not txt:
                return False, "empty"
            t = txt.lower().strip()
            a = gt["answer"]
            if qt in ("pair_distance", "view_angle", "free_space"):
                nums = re.findall(r'[\d.]+', txt)
                if not nums:
                    return False, "no_num"
                try:
                    pv = float(nums[-1])
                    tol = {"pair_distance": max(0.02, abs(a) * 0.05), "view_angle": 5.0, "free_space": max(3.0, abs(a) * 0.1)}[qt]
                    return abs(pv - a) <= tol, f"pred={pv} gt={a}"
                except:
                    return False, "parse_err"
            elif qt == "obstruction":
                if isinstance(a, list) and a == ["none"]:
                    return ("none" in t or "no object" in t or "no obstruction" in t), "none"
                found = sum(1 for x in a if x.lower() in t)
                return found >= max(1, len(a) // 2), f"found={found}/{len(a)}"
            elif qt == "placement":
                if isinstance(a, dict):
                    correct = sum(1 for lab, zone in a.items() if zone.lower() in t)
                    return correct >= max(1, len(a) // 2), f"ok={correct}/{len(a)}"
                return False, "bad_gt"
            elif qt == "repositioning":
                return a.lower() in t, f"gt={a}"
            return False, "unknown"

        items = []
        for rt in ["kitchen", "living_room", "bedroom", "hssd"]:
            rdir = os.path.join(layout_dir, rt)
            files = sorted([f for f in os.listdir(rdir) if f.endswith(".json")])[:sample]
            for fname in files:
                data = json.load(open(os.path.join(rdir, fname)))
                if perturb > 0:
                    d2 = copy.deepcopy(data)
                    rng = random.Random(abs(hash(d2.get("layout_id", 0))) + 12345)
                    def j(pts):
                        return [{"x": round(p["x"] + rng.gauss(0, perturb), 4),
                                 "y": round(p["y"] + rng.gauss(0, perturb), 4)} for p in pts]
                    if "room_boundary" in d2:
                        d2["room_boundary"] = j(d2["room_boundary"])
                    for o in d2.get("objects", []):
                        o["points"] = j(o["points"])
                    data = d2
                for qt in QTYPES:
                    gt = gen_gt(qt, data)
                    if gt:
                        items.append({"rt": rt, "qt": qt, "cat": CATS[qt], "lid": data.get("layout_id"), "gt": gt, "data": data})

        print(f"Benchmark: {len(items)} items", flush=True)

        sys_msg = "You are a spatial reasoning assistant. Given a structured floor layout, answer the spatial question. Reply with just the answer."
        results = []
        t0 = time.time()
        for idx, it in enumerate(items):
            user = make_prompt(it["qt"], it["gt"], it["data"], fmt)
            try:
                txt = self._query([{"role": "system", "content": sys_msg}, {"role": "user", "content": user}])
            except Exception as e:
                txt = f"__ERROR__:{e}"
            ok, reason = score(it["qt"], it["gt"], txt)
            results.append({"rt": it["rt"], "qt": it["qt"], "cat": it["cat"], "lid": it["lid"], "ok": bool(ok), "reason": reason})
            if (idx + 1) % 20 == 0:
                acc = sum(r["ok"] for r in results) / len(results)
                print(f"  {idx + 1}/{len(items)}  {time.time() - t0:.0f}s  acc={acc:.3f}", flush=True)

        overall = sum(r["ok"] for r in results) / max(1, len(results))
        by_cat = defaultdict(lambda: {"c": 0, "n": 0})
        by_rt = defaultdict(lambda: {"c": 0, "n": 0})
        for r in results:
            by_cat[r["cat"]]["c"] += int(r["ok"])
            by_cat[r["cat"]]["n"] += 1
            by_rt[r["rt"]]["c"] += int(r["ok"])
            by_rt[r["rt"]]["n"] += 1

        summary = {
            "model": self.model_name, "fmt": fmt, "perturb": perturb, "sample": sample,
            "n_items": len(results), "overall_acc": round(overall, 4), "elapsed_sec": round(time.time() - t0, 1),
            "by_cat": {k: {"acc": round(v["c"] / v["n"], 4), "n": v["n"]} for k, v in by_cat.items()},
            "by_rt": {k: {"acc": round(v["c"] / v["n"], 4), "n": v["n"]} for k, v in by_rt.items()},
        }
        print(json.dumps(summary, indent=2), flush=True)
        return summary


@app.local_entrypoint()
def main():
    results_dir = "results"
    os.makedirs(results_dir, exist_ok=True)

    # JSON baseline
    print(f"\n=== {MODEL} (JSON) ===", flush=True)
    e = FPQAEval(model_name=MODEL)
    r = e.evaluate.remote(sample=SAMPLE, fmt="json", perturb=0.0)
    path = os.path.join(results_dir, "Qwen2.5-7B-Instruct_modal.json")
    with open(path, "w") as f:
        json.dump(r, f, indent=2)
    print(f"  -> acc={r['overall_acc']}", flush=True)

    # XML ablation
    print("\n=== XML ABLATION ===", flush=True)
    r_xml = e.evaluate.remote(sample=SAMPLE, fmt="xml", perturb=0.0)
    with open(os.path.join(results_dir, "Qwen2.5-7B-Instruct_modal_xml.json"), "w") as f:
        json.dump(r_xml, f, indent=2)

    # Perturbation
    print("\n=== PERTURBATION ===", flush=True)
    r_pert = e.evaluate.remote(sample=SAMPLE, fmt="json", perturb=0.1)
    with open(os.path.join(results_dir, "Qwen2.5-7B-Instruct_modal_perturb.json"), "w") as f:
        json.dump(r_pert, f, indent=2)

    print("\n=== DONE ===", flush=True)