atropos/environments/eval_environments/vision_evals/hallusionbench_environment.py
2026-01-23 00:49:51 +00:00

147 lines
5 KiB
Python

"""HallusionBench evaluation environment."""
import asyncio
import base64
import io
import re
from typing import List, Optional, Tuple
from datasets import load_dataset
from environments.eval_environments.eval import EvalBase, eval_runner
from PIL import Image
from atroposlib.envs.server_handling.server_manager import ServerManager
class HallusionBench(EvalBase):
"""HallusionBench evaluation - visual hallucination benchmark."""
def setup_data(self) -> list:
# HallusionBench has 'image' and 'non_image' splits
split = getattr(self, "split", "image")
try:
dataset = load_dataset("lmms-lab/HallusionBench", split=split)
print(f"Loaded {len(dataset)} examples from HallusionBench ({split})")
return list(dataset)
except Exception as e:
print(f"Warning: Could not load HallusionBench: {e}")
try:
# Try combining both splits
all_data = []
for s in ["image", "non_image"]:
try:
ds = load_dataset("lmms-lab/HallusionBench", split=s)
all_data.extend(list(ds))
except Exception:
pass
if all_data:
print(
f"Loaded {len(all_data)} examples from HallusionBench (combined)"
)
return all_data
raise ValueError(f"Could not load HallusionBench dataset: {e}")
except Exception:
raise ValueError(f"Could not load HallusionBench dataset: {e}")
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) -> Optional[str]:
for key in ["image", "decoded_image"]:
if key in item and item[key] is not None:
if isinstance(item[key], Image.Image):
return self.encode_image(item[key])
return None
def build_messages(self, item: dict) -> List[dict]:
image_base64 = self.get_image_base64(item)
question = item.get("question", "")
prompt = f"{question}\n\nPlease answer yes or no."
content = []
if image_base64:
content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{image_base64}"},
}
)
content.append({"type": "text", "text": prompt})
return [{"role": "user", "content": content}]
def extract_yorn(self, response: str) -> str:
"""Extract Yes/No from response."""
response_lower = response.lower().strip()
if response_lower.startswith("yes"):
return "Yes"
if response_lower.startswith("no"):
return "No"
yes_patterns = [r"\byes\b", r"\btrue\b", r"\bcorrect\b"]
no_patterns = [r"\bno\b", r"\bfalse\b", r"\bincorrect\b"]
for pattern in yes_patterns:
if re.search(pattern, response_lower):
return "Yes"
for pattern in no_patterns:
if re.search(pattern, response_lower):
return "No"
return "Unknown"
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 not response:
return {"accuracy": 0.0}, {"error": "Empty response"}
answer = data_item.get("answer", data_item.get("gt_answer", ""))
extracted = self.extract_yorn(response)
answer_norm = str(answer).strip().lower()
if answer_norm in ["yes", "true", "1"]:
answer_norm = "Yes"
elif answer_norm in ["no", "false", "0"]:
answer_norm = "No"
else:
answer_norm = str(answer).strip()
correct = extracted == answer_norm
sample = {
"id": data_item.get("index", data_item.get("id", "")),
"question": data_item.get("question", "")[:200],
"category": data_item.get("category", data_item.get("subcategory", "")),
"answer": answer_norm,
"prediction": extracted,
"raw_response": response[:200],
"correct": correct,
}
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(HallusionBench(split="test", temperature=0.0, max_tokens=64))
)