evals moved + readme

environments/eval_environments/vision_evals/visulogic_environment.py

@@ -0,0 +1,182 @@
+import asyncio
+import base64
+import io
+import json
+import zipfile
+from pathlib import Path
+from typing import List, Tuple
+
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from atroposlib.envs.server_handling.server_manager import ServerManager
+from environments.eval_environments.eval import EvalBase, eval_runner
+
+DEFAULT_DATA_DIR = Path.home() / ".cache" / "visulogic_hf"
+
+
+class VisuLogic(EvalBase):
+    TAGS = [
+        "Quantitative Reasoning",
+        "Spatial Reasoning",
+        "Positional Reasoning",
+        "Attribute Reasoning",
+        "Stylistic Reasoning",
+        "Other",
+    ]
+
+    def _download_data(self, data_dir: Path) -> None:
+        jsonl_path = data_dir / "data.jsonl"
+        images_dir = data_dir / "images"
+
+        if jsonl_path.exists() and images_dir.exists():
+            return
+
+        print(f"Downloading VisuLogic dataset to {data_dir}...")
+        data_dir.mkdir(parents=True, exist_ok=True)
+
+        # Download data.jsonl
+        hf_hub_download(
+            repo_id="VisuLogic/VisuLogic",
+            filename="data.jsonl",
+            repo_type="dataset",
+            local_dir=data_dir,
+        )
+
+        # Download and extract images.zip
+        images_zip_path = hf_hub_download(
+            repo_id="VisuLogic/VisuLogic",
+            filename="images.zip",
+            repo_type="dataset",
+            local_dir=data_dir,
+        )
+
+        print("Extracting images...")
+        with zipfile.ZipFile(images_zip_path, "r") as zip_ref:
+            zip_ref.extractall(data_dir)
+
+        print("Download complete!")
+
+    def setup_data(self) -> list:
+        """
+        Load and return dataset as a list.
+
+        Auto-downloads the VisuLogic dataset if data_path is not specified
+        or doesn't exist.
+        """
+        data_path = getattr(self, "data_path", None)
+
+        if data_path is None:
+            data_dir = DEFAULT_DATA_DIR
+            self._download_data(data_dir)
+            jsonl_path = data_dir / "data.jsonl"
+            self.images_base = str(data_dir)
+        else:
+            data_dir = Path(data_path)
+            jsonl_path = data_dir / "data.jsonl"
+            self.images_base = str(data_dir)
+
+            if not jsonl_path.exists():
+                raise FileNotFoundError(
+                    f"Dataset not found at {jsonl_path}. "
+                    "Remove data_path argument to auto-download."
+                )
+
+        dataset = []
+        with open(jsonl_path, "r", encoding="utf-8") as f:
+            for line in f:
+                item = json.loads(line.strip())
+                dataset.append(item)
+
+        print(f"Loaded {len(dataset)} examples from VisuLogic")
+        return dataset
+
+    def encode_image(self, pil_image: Image.Image) -> str:
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format="PNG")
+        return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+    def get_image_base64(self, item: dict) -> str:
+        image_path = item.get("image_path", "")
+        full_path = Path(self.images_base) / image_path
+        if full_path.exists():
+            with Image.open(full_path) as img:
+                return self.encode_image(img)
+        raise ValueError(f"Could not find image at {full_path}")
+
+    def build_messages(self, item: dict) -> List[dict]:
+        image_base64 = self.get_image_base64(item)
+        question = item.get("question", "")
+
+        prompt = f"""{question}
+
+Answer with only the letter (A, B, C, or D)."""
+
+        return [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/png;base64,{image_base64}"},
+                    },
+                    {"type": "text", "text": prompt},
+                ],
+            }
+        ]
+
+    def extract_answer(self, response: str) -> str:
+        response = response.strip().upper()
+
+        for char in reversed(response):
+            if char in "ABCD":
+                return char
+
+        return ""
+
+    def score(self, prediction: str, answer: str) -> bool:
+        if not prediction:
+            return False
+        return prediction.upper() == answer.upper()
+
+    async def run_item(self, server: ServerManager, data_item: dict) -> Tuple[dict, dict]:
+        try:
+            messages = self.build_messages(data_item)
+
+            completion = await self.chat_completion(server, messages)
+
+            if not completion.choices:
+                return {"accuracy": 0.0}, {"error": "Empty response"}
+
+            message = completion.choices[0].message
+            response = message.content or ""
+            if hasattr(message, "reasoning") and message.reasoning and not response:
+                response = message.reasoning
+            if not response and hasattr(message, "model_extra"):
+                reasoning = message.model_extra.get("reasoning", "")
+                if reasoning:
+                    response = reasoning
+
+            if not response:
+                return {"accuracy": 0.0}, {"error": "Empty response"}
+
+            extracted = self.extract_answer(response)
+            answer = data_item.get("label", "")
+            correct = self.score(extracted, answer)
+
+            sample = {
+                "question": data_item.get("question", ""),
+                "answer": answer,
+                "prediction": extracted,
+                "correct": correct,
+                "tag": data_item.get("tag", ""),
+            }
+
+            return {"accuracy": 1.0 if correct else 0.0}, sample
+
+        except Exception as e:
+            return {"accuracy": 0.0}, {"error": str(e)}
+
+
+if __name__ == "__main__":
+    asyncio.run(eval_runner(VisuLogic(temperature=0.0, max_tokens=256)))