mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
223 lines
7.7 KiB
Python
223 lines
7.7 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
EVAL_DIR = Path(__file__).resolve().parent.parent # eval/
|
||
if str(EVAL_DIR) not in sys.path:
|
||
sys.path.insert(0, str(EVAL_DIR))
|
||
|
||
from common.io import require_file # noqa: E402
|
||
from common.paths import DATA_DIR # noqa: E402
|
||
|
||
# The official three-part answer format from hle_eval
|
||
HLE_ANSWER_FORMAT = (
|
||
"{question}\n"
|
||
"Your response should be in the following format:\n"
|
||
"Explanation: {{your explanation for your final answer}}\n"
|
||
"Answer: {{your succinct, final answer}}\n"
|
||
"Confidence: {{your confidence score between 0% and 100% for your answer}}"
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Problem:
|
||
task_id: str
|
||
question: str
|
||
answer: str
|
||
category: str = "unknown"
|
||
has_image: bool = False
|
||
image_paths: tuple[str, ...] = ()
|
||
image_sizes: tuple[tuple[int, int], ...] = ()
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Data
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
_IMAGE_MAGIC = (b"\xff\xd8\xff", b"\x89PNG", b"GIF8", b"RIFF", b"BM")
|
||
|
||
|
||
def _resolve_images(row: dict[str, Any]) -> tuple[tuple[str, ...], tuple[tuple[int, int], ...]]:
|
||
"""Resolve a row's image into (local file path, (W, H)).
|
||
|
||
Three representations are supported:
|
||
1. multi_modal_data.image_urls + image_sizes (the full-dataset JSONL uses this form)
|
||
2. an `image` field holding a local path
|
||
3. an `image` field holding base64 (common in the official cais/hle parquet), decoded and
|
||
cached under data/hle_images/
|
||
"""
|
||
import base64
|
||
|
||
multi_modal = row.get("multi_modal_data") or {}
|
||
urls = [str(url) for url in (multi_modal.get("image_urls") or [])]
|
||
sizes = [(int(size[0]), int(size[1])) for size in (multi_modal.get("image_sizes") or []) if len(size) >= 2]
|
||
|
||
raw = row.get("image")
|
||
if not urls and raw:
|
||
if isinstance(raw, (bytes, bytearray)):
|
||
urls = [_write_cached_image(row, bytes(raw), 0)]
|
||
else:
|
||
text = str(raw).strip()
|
||
if os.path.isfile(text):
|
||
urls = [text]
|
||
elif text:
|
||
try:
|
||
data = base64.b64decode(text, validate=False)
|
||
except Exception:
|
||
data = b""
|
||
if any(data.startswith(magic) for magic in _IMAGE_MAGIC):
|
||
urls = [_write_cached_image(row, data, 0)]
|
||
|
||
if urls and not sizes:
|
||
from PIL import Image # PIL is only needed to recover the dimensions from a file
|
||
|
||
sizes = [Image.open(path).size for path in urls]
|
||
return tuple(urls), tuple(sizes)
|
||
|
||
|
||
def _write_cached_image(row: dict[str, Any], data: bytes, index: int) -> str:
|
||
cache = DATA_DIR / "hle_images"
|
||
cache.mkdir(parents=True, exist_ok=True)
|
||
stem = str(row.get("task_id") or row.get("id") or "hle")
|
||
path = cache / f"{stem}_{index}.img"
|
||
if not path.is_file() or path.stat().st_size != len(data):
|
||
path.write_bytes(data)
|
||
return str(path)
|
||
|
||
|
||
def _has_image(row: dict[str, Any]) -> bool:
|
||
"""Accept several common representations: an explicit modality field, multi_modal_data.image_urls, or a single image field."""
|
||
if str(row.get("modality") or "").lower() == "image":
|
||
return True
|
||
multi_modal = row.get("multi_modal_data") or {}
|
||
if multi_modal.get("image_urls"):
|
||
return True
|
||
img = row.get("image")
|
||
if img is None:
|
||
return False
|
||
if isinstance(img, (bytes, bytearray)):
|
||
return len(img) > 0
|
||
return bool(str(img).strip())
|
||
|
||
|
||
def load(data_file: str | os.PathLike) -> list[Problem]:
|
||
path = Path(data_file)
|
||
require_file(path, "HLE download: https://huggingface.co/datasets/cais/hle (gated; accept the terms on Hugging Face first)")
|
||
|
||
problems: list[Problem] = []
|
||
with open(path) as f:
|
||
for line_no, line in enumerate(f):
|
||
line = line.strip()
|
||
if not line:
|
||
continue
|
||
row = json.loads(line)
|
||
has_image = _has_image(row)
|
||
image_paths, image_sizes = _resolve_images(row) if has_image else ((), ())
|
||
problems.append(
|
||
Problem(
|
||
task_id=str(row.get("task_id") or row.get("id") or f"hle_{line_no:04d}"),
|
||
question=str(row["question"]),
|
||
answer=str(row["answer"]),
|
||
category=str(row.get("category") or row.get("raw_subject") or "unknown"),
|
||
has_image=has_image,
|
||
image_paths=image_paths,
|
||
image_sizes=image_sizes,
|
||
)
|
||
)
|
||
assert len({p.task_id for p in problems}) == len(problems), "duplicate task_id"
|
||
return problems
|
||
|
||
|
||
def build_question(problem: Problem) -> str:
|
||
return HLE_ANSWER_FORMAT.format(question=problem.question.strip())
|
||
|
||
|
||
def build_messages(problem: Problem) -> list[dict[str, Any]]:
|
||
return [{"role": "user", "content": build_question(problem)}]
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Answer extraction and (rule-based) scoring
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
_ANSWER_LINE_RE = re.compile(r"(?im)^\s*(?:\*\*)?answer(?:\*\*)?\s*[:：]\s*(.+?)\s*$")
|
||
|
||
|
||
def extract_answer(text: str) -> str | None:
|
||
"""Return the content of the last `Answer:` line (the official three-part format)."""
|
||
if not text:
|
||
return None
|
||
matches = _ANSWER_LINE_RE.findall(text)
|
||
if not matches:
|
||
return None
|
||
ans = matches[-1].strip()
|
||
ans = ans.strip("*").strip()
|
||
ans = re.sub(r"^\$+|\$+$", "", ans).strip()
|
||
return ans or None
|
||
|
||
|
||
def normalize_answer(ans: str | None) -> str | None:
|
||
if ans is None:
|
||
return None
|
||
text = ans.strip()
|
||
text = text.strip("$").strip()
|
||
text = text.replace("\\left", "").replace("\\right", "")
|
||
text = text.replace("\\!", "").replace("\\,", "").replace("\\;", "")
|
||
text = text.replace("{,}", ",") if "{,}" in text else text
|
||
text = re.sub(r"\s+", " ", text).strip().lower()
|
||
text = text.rstrip(".").strip()
|
||
if re.fullmatch(r"-?\d+", text):
|
||
return str(int(text))
|
||
return text
|
||
|
||
|
||
def is_correct(pred: str | None, gold: str) -> bool:
|
||
"""Rule-based scoring: normalized string equality. **Smoke tests only**; report LLM-judge numbers."""
|
||
pred_norm, gold_norm = normalize_answer(pred), normalize_answer(gold)
|
||
return pred_norm is not None and gold_norm is not None and pred_norm == gold_norm
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# Entry points used by run_eval.py / score.py
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
|
||
def build_records(
|
||
problems: list[Problem],
|
||
*,
|
||
limit: int | None = None,
|
||
) -> list[dict[str, Any]]:
|
||
records = []
|
||
for problem in problems:
|
||
records.append(
|
||
{
|
||
"task_id": problem.task_id,
|
||
"question": problem.question,
|
||
"answer": problem.answer,
|
||
"category": problem.category,
|
||
"modality": "image" if problem.has_image else "text",
|
||
"image_paths": list(problem.image_paths),
|
||
"image_sizes": [list(size) for size in problem.image_sizes],
|
||
"judge_question": problem.question,
|
||
"messages": build_messages(problem),
|
||
}
|
||
)
|
||
if limit is not None and len(records) >= limit:
|
||
break
|
||
return records
|
||
|
||
|
||
def score_rollout(rollout: dict[str, Any], record: dict[str, Any]) -> dict[str, Any]:
|
||
row = dict(rollout)
|
||
row["pred"] = extract_answer(rollout.get("raw_text", ""))
|
||
row["gold"] = record["answer"]
|
||
row["correct"] = is_correct(row["pred"], record["answer"])
|
||
return row
|