571 lines
20 KiB
Python
571 lines
20 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Recursive inference engine for stateless recursive math-solving models."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
|
||
|
|
TAG_RE = re.compile(r"</?(problem|action|parallel|path)>", re.IGNORECASE)
|
||
|
|
NUMBER_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?|-?\d+/\d+")
|
||
|
|
FINAL_ANSWER_LINE_RE = re.compile(
|
||
|
|
r"(?im)^[ \t]*Final[ \t]+answer[ \t]*:[ \t]*(?P<answer>[^\n]+)$"
|
||
|
|
)
|
||
|
|
DEFAULT_PREFIX_PROMPT_FILE = Path("prompts") / "recursive_sft_prefix_prompt.txt"
|
||
|
|
FORCED_ROOT_OUTPUT_PREFIXES = {
|
||
|
|
"SOLVE_DIRECTLY": "<action>\nSOLVE_DIRECTLY\n",
|
||
|
|
"DECOMPOSE": "<action>\nDECOMPOSE\n</action>\n<parallel>\n",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class ModelCall:
|
||
|
|
call_id: int
|
||
|
|
depth: int
|
||
|
|
role: str
|
||
|
|
problem: str
|
||
|
|
input_text: str
|
||
|
|
output_text: str
|
||
|
|
input_tokens: int
|
||
|
|
output_tokens: int
|
||
|
|
action_type: str | None = None
|
||
|
|
parse_error: str | None = None
|
||
|
|
forced_output_prefix: str = ""
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class RolloutResult:
|
||
|
|
problem: str
|
||
|
|
final_answer: str | None
|
||
|
|
full_trace: str
|
||
|
|
calls: list[ModelCall] = field(default_factory=list)
|
||
|
|
parse_errors: list[str] = field(default_factory=list)
|
||
|
|
critical_path_tokens: int = 0
|
||
|
|
stopped_reason: str = "completed"
|
||
|
|
|
||
|
|
@property
|
||
|
|
def parsed_ok(self) -> bool:
|
||
|
|
return not self.parse_errors
|
||
|
|
|
||
|
|
@property
|
||
|
|
def model_call_count(self) -> int:
|
||
|
|
return len(self.calls)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class NodeResult:
|
||
|
|
final_answer: str | None
|
||
|
|
trace: str
|
||
|
|
critical_path_tokens: int
|
||
|
|
parse_errors: list[str]
|
||
|
|
stopped_reason: str = "completed"
|
||
|
|
|
||
|
|
|
||
|
|
def load_prefix_template(path: str | Path | None = None) -> str:
|
||
|
|
prompt_path = Path(path) if path is not None else DEFAULT_PREFIX_PROMPT_FILE
|
||
|
|
return prompt_path.read_text(encoding="utf-8")
|
||
|
|
|
||
|
|
|
||
|
|
def format_model_input(problem: str, prefix_template: str) -> str:
|
||
|
|
return prefix_template.strip().replace("{PROBLEM}", problem.strip())
|
||
|
|
|
||
|
|
|
||
|
|
def token_count(tokenizer: Any, text: str) -> int:
|
||
|
|
return len(tokenizer.encode(text, add_special_tokens=False))
|
||
|
|
|
||
|
|
|
||
|
|
def extract_blocks(text: str, tag: str) -> list[str]:
|
||
|
|
return [
|
||
|
|
match.group(1).strip()
|
||
|
|
for match in re.finditer(rf"<{tag}>(.*?)</{tag}>", text, flags=re.DOTALL | re.IGNORECASE)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_answer(value: Any) -> str | None:
|
||
|
|
if value is None:
|
||
|
|
return None
|
||
|
|
text = str(value).strip()
|
||
|
|
if not text:
|
||
|
|
return None
|
||
|
|
if "####" in text:
|
||
|
|
text = text.split("####")[-1].strip()
|
||
|
|
boxed = re.findall(r"\\boxed\{([^{}]+)\}", text)
|
||
|
|
if boxed:
|
||
|
|
text = boxed[-1]
|
||
|
|
final = re.findall(r"(?i)Final\s+answer\s*:\s*([^\n<]+)", text)
|
||
|
|
if final:
|
||
|
|
text = final[-1]
|
||
|
|
text = text.replace("$", "").replace(",", "")
|
||
|
|
text = re.sub(r"\\(?:text|mathrm)\{([^{}]+)\}", r"\1", text)
|
||
|
|
text = re.sub(r"\s+", "", text)
|
||
|
|
text = text.rstrip(".")
|
||
|
|
return text or None
|
||
|
|
|
||
|
|
|
||
|
|
def extract_final_answer(text: str) -> str | None:
|
||
|
|
patterns = [
|
||
|
|
r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
|
||
|
|
r"(?i)answer\s+is\s*([^\n<]+)",
|
||
|
|
r"\\boxed\{([^{}]+)\}",
|
||
|
|
]
|
||
|
|
for pattern in patterns:
|
||
|
|
matches = re.findall(pattern, text)
|
||
|
|
if matches:
|
||
|
|
return normalize_answer(matches[-1])
|
||
|
|
matches = NUMBER_RE.findall(text)
|
||
|
|
if matches:
|
||
|
|
return normalize_answer(matches[-1])
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def answers_match(pred: str | None, gt: str | None) -> bool:
|
||
|
|
pred_norm = normalize_answer(pred)
|
||
|
|
gt_norm = normalize_answer(gt)
|
||
|
|
if pred_norm is None or gt_norm is None:
|
||
|
|
return False
|
||
|
|
|
||
|
|
if pred_norm == gt_norm:
|
||
|
|
return True
|
||
|
|
|
||
|
|
try:
|
||
|
|
from math_verify import LatexExtractionConfig, parse, verify
|
||
|
|
|
||
|
|
pred_parsed = parse(pred_norm, extraction_config=[LatexExtractionConfig()])
|
||
|
|
gt_parsed = parse(gt_norm, extraction_config=[LatexExtractionConfig()])
|
||
|
|
if pred_parsed and gt_parsed and verify(gt_parsed, pred_parsed):
|
||
|
|
return True
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
|
||
|
|
return pred_norm.lower() == gt_norm.lower()
|
||
|
|
|
||
|
|
|
||
|
|
def find_matching_tag(text: str, tag: str, open_start: int, open_end: int) -> tuple[int, int] | None:
|
||
|
|
depth = 1
|
||
|
|
for match in TAG_RE.finditer(text, open_end):
|
||
|
|
token = match.group(0).lower()
|
||
|
|
name = match.group(1).lower()
|
||
|
|
if name != tag:
|
||
|
|
continue
|
||
|
|
if token.startswith("</"):
|
||
|
|
depth -= 1
|
||
|
|
if depth == 0:
|
||
|
|
return match.start(), match.end()
|
||
|
|
else:
|
||
|
|
depth += 1
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def parse_model_output(output: str) -> tuple[str, list[str]]:
|
||
|
|
action_match = re.search(r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE)
|
||
|
|
if not action_match:
|
||
|
|
raise ValueError("missing <action>...</action> block")
|
||
|
|
|
||
|
|
action_body = action_match.group(1).strip()
|
||
|
|
first_line = action_body.splitlines()[0].strip().upper() if action_body else ""
|
||
|
|
if first_line.startswith("SOLVE_DIRECTLY"):
|
||
|
|
return "SOLVE_DIRECTLY", []
|
||
|
|
if first_line.startswith("AGGREGATE"):
|
||
|
|
return "AGGREGATE", []
|
||
|
|
if not first_line.startswith("DECOMPOSE"):
|
||
|
|
raise ValueError(f"unknown action type in first action line: {first_line!r}")
|
||
|
|
|
||
|
|
next_nonspace = re.search(r"\S", output[action_match.end() :])
|
||
|
|
if not next_nonspace:
|
||
|
|
raise ValueError("DECOMPOSE action is not followed by <parallel>")
|
||
|
|
|
||
|
|
parallel_start = action_match.end() + next_nonspace.start()
|
||
|
|
if not output[parallel_start:].lower().startswith("<parallel>"):
|
||
|
|
raise ValueError("DECOMPOSE action is not immediately followed by <parallel>")
|
||
|
|
|
||
|
|
parallel_open_end = parallel_start + len("<parallel>")
|
||
|
|
close = find_matching_tag(output, "parallel", parallel_start, parallel_open_end)
|
||
|
|
if close is None:
|
||
|
|
raise ValueError("unclosed <parallel> block")
|
||
|
|
parallel_body = output[parallel_open_end : close[0]]
|
||
|
|
paths = extract_blocks(f"<parallel>{parallel_body}</parallel>", "path")
|
||
|
|
if len(paths) < 2:
|
||
|
|
raise ValueError(f"DECOMPOSE parallel block has {len(paths)} path(s), expected at least 2")
|
||
|
|
if any(not path for path in paths):
|
||
|
|
raise ValueError("DECOMPOSE parallel block contains an empty path")
|
||
|
|
return "DECOMPOSE", paths
|
||
|
|
|
||
|
|
|
||
|
|
def validate_action_protocol(output: str, action_type: str, role: str) -> None:
|
||
|
|
"""Enforce role-specific actions and explicit answers for terminal calls."""
|
||
|
|
normalized_role = role.strip().lower()
|
||
|
|
if normalized_role == "root":
|
||
|
|
allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
|
||
|
|
elif normalized_role == "aggregate":
|
||
|
|
allowed_actions = {"AGGREGATE"}
|
||
|
|
elif normalized_role == "child" or normalized_role.startswith("path_"):
|
||
|
|
allowed_actions = {"DECOMPOSE", "SOLVE_DIRECTLY"}
|
||
|
|
else:
|
||
|
|
raise ValueError(f"unknown call role: {role!r}")
|
||
|
|
|
||
|
|
if action_type not in allowed_actions:
|
||
|
|
expected = " or ".join(sorted(allowed_actions))
|
||
|
|
raise ValueError(
|
||
|
|
f"{normalized_role} call returned {action_type}, expected {expected}"
|
||
|
|
)
|
||
|
|
|
||
|
|
if action_type not in {"SOLVE_DIRECTLY", "AGGREGATE"}:
|
||
|
|
return
|
||
|
|
|
||
|
|
action_match = re.search(
|
||
|
|
r"<action>(.*?)</action>", output, flags=re.DOTALL | re.IGNORECASE
|
||
|
|
)
|
||
|
|
action_body = action_match.group(1) if action_match else ""
|
||
|
|
final_answer = FINAL_ANSWER_LINE_RE.search(action_body)
|
||
|
|
if final_answer is None or not final_answer.group("answer").strip():
|
||
|
|
raise ValueError(f"{action_type} action is missing a non-empty Final answer: line")
|
||
|
|
|
||
|
|
|
||
|
|
def extract_final_answer_display(text: str) -> str | None:
|
||
|
|
patterns = [
|
||
|
|
r"(?i)Final\s+answer\s*:\s*([^\n<]+)",
|
||
|
|
r"\\boxed\{([^{}]+)\}",
|
||
|
|
]
|
||
|
|
for pattern in patterns:
|
||
|
|
matches = re.findall(pattern, text)
|
||
|
|
if matches:
|
||
|
|
return matches[-1].strip().rstrip(".")
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def child_result_summary(result: NodeResult) -> str:
|
||
|
|
display = extract_final_answer_display(result.trace)
|
||
|
|
if display is not None:
|
||
|
|
return display
|
||
|
|
if result.final_answer is not None:
|
||
|
|
return str(result.final_answer)
|
||
|
|
fallback = extract_final_answer(result.trace)
|
||
|
|
if fallback is not None:
|
||
|
|
return str(fallback)
|
||
|
|
return "PARSE_ERROR" if result.parse_errors else "UNKNOWN"
|
||
|
|
|
||
|
|
|
||
|
|
def build_aggregation_problem(parent_problem: str, paths: list[str], child_results: list[NodeResult]) -> str:
|
||
|
|
lines = [
|
||
|
|
f"The original problem is: {parent_problem.strip()}",
|
||
|
|
"The paths spawned were:",
|
||
|
|
]
|
||
|
|
for idx, (path, result) in enumerate(zip(paths, child_results), start=1):
|
||
|
|
summary = child_result_summary(result)
|
||
|
|
lines.append(f"{idx}. {path.strip()} (Result: {summary}).")
|
||
|
|
lines.append("Aggregate these results to find the answer to the original problem.")
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
|
||
|
|
class RecursiveEngine:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
model: Any,
|
||
|
|
tokenizer: Any,
|
||
|
|
prefix_template: str,
|
||
|
|
max_new_tokens: int = 768,
|
||
|
|
temperature: float = 0.0,
|
||
|
|
top_p: float = 1.0,
|
||
|
|
top_k: int | None = None,
|
||
|
|
max_depth: int = 3,
|
||
|
|
max_calls: int = 10,
|
||
|
|
force_root_action: str | None = None,
|
||
|
|
) -> None:
|
||
|
|
self.model = model
|
||
|
|
self.tokenizer = tokenizer
|
||
|
|
self.prefix_template = prefix_template
|
||
|
|
self.max_new_tokens = max_new_tokens
|
||
|
|
self.temperature = temperature
|
||
|
|
self.top_p = top_p
|
||
|
|
self.top_k = top_k
|
||
|
|
self.max_depth = max_depth
|
||
|
|
self.max_calls = max_calls
|
||
|
|
normalized_force = force_root_action.upper() if force_root_action else None
|
||
|
|
if normalized_force not in {None, *FORCED_ROOT_OUTPUT_PREFIXES}:
|
||
|
|
raise ValueError(
|
||
|
|
"force_root_action must be SOLVE_DIRECTLY, DECOMPOSE, or None"
|
||
|
|
)
|
||
|
|
self.force_root_action = normalized_force
|
||
|
|
self.calls: list[ModelCall] = []
|
||
|
|
|
||
|
|
def generate_once(self, input_text: str, *, output_prefix: str = "") -> str:
|
||
|
|
import torch
|
||
|
|
|
||
|
|
generation_input = input_text + output_prefix
|
||
|
|
inputs = self.tokenizer(generation_input, return_tensors="pt")
|
||
|
|
device = next(self.model.parameters()).device
|
||
|
|
inputs = {key: value.to(device) for key, value in inputs.items()}
|
||
|
|
generation_kwargs = {
|
||
|
|
"max_new_tokens": self.max_new_tokens,
|
||
|
|
"do_sample": self.temperature > 0,
|
||
|
|
"top_p": self.top_p,
|
||
|
|
"use_cache": True,
|
||
|
|
"pad_token_id": self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
|
||
|
|
"eos_token_id": self.tokenizer.eos_token_id,
|
||
|
|
}
|
||
|
|
if self.temperature > 0:
|
||
|
|
generation_kwargs["temperature"] = self.temperature
|
||
|
|
if self.top_k is not None:
|
||
|
|
generation_kwargs["top_k"] = self.top_k
|
||
|
|
with torch.no_grad():
|
||
|
|
generated = self.model.generate(
|
||
|
|
**inputs,
|
||
|
|
**generation_kwargs,
|
||
|
|
)
|
||
|
|
new_tokens = generated[0, inputs["input_ids"].shape[1] :]
|
||
|
|
continuation = self.tokenizer.decode(new_tokens, skip_special_tokens=False).strip()
|
||
|
|
return output_prefix + continuation
|
||
|
|
|
||
|
|
def rollout(self, problem: str) -> RolloutResult:
|
||
|
|
self.calls = []
|
||
|
|
root = self._solve_node(problem=problem, depth=0, role="root")
|
||
|
|
return RolloutResult(
|
||
|
|
problem=problem,
|
||
|
|
final_answer=root.final_answer,
|
||
|
|
full_trace=root.trace,
|
||
|
|
calls=self.calls,
|
||
|
|
parse_errors=root.parse_errors,
|
||
|
|
critical_path_tokens=root.critical_path_tokens,
|
||
|
|
stopped_reason=root.stopped_reason,
|
||
|
|
)
|
||
|
|
|
||
|
|
def _solve_node(self, *, problem: str, depth: int, role: str) -> NodeResult:
|
||
|
|
if len(self.calls) >= self.max_calls:
|
||
|
|
return NodeResult(
|
||
|
|
final_answer=None,
|
||
|
|
trace="",
|
||
|
|
critical_path_tokens=0,
|
||
|
|
parse_errors=[f"max model calls exceeded ({self.max_calls})"],
|
||
|
|
stopped_reason="max_calls",
|
||
|
|
)
|
||
|
|
if depth > self.max_depth:
|
||
|
|
return NodeResult(
|
||
|
|
final_answer=None,
|
||
|
|
trace="",
|
||
|
|
critical_path_tokens=0,
|
||
|
|
parse_errors=[f"max recursion depth exceeded ({self.max_depth})"],
|
||
|
|
stopped_reason="max_depth",
|
||
|
|
)
|
||
|
|
|
||
|
|
input_text = format_model_input(problem, self.prefix_template)
|
||
|
|
forced_output_prefix = ""
|
||
|
|
if depth == 0 and role == "root" and self.force_root_action is not None:
|
||
|
|
forced_output_prefix = FORCED_ROOT_OUTPUT_PREFIXES[self.force_root_action]
|
||
|
|
output_text = self.generate_once(
|
||
|
|
input_text,
|
||
|
|
output_prefix=forced_output_prefix,
|
||
|
|
)
|
||
|
|
generated_output = (
|
||
|
|
output_text[len(forced_output_prefix) :]
|
||
|
|
if forced_output_prefix and output_text.startswith(forced_output_prefix)
|
||
|
|
else output_text
|
||
|
|
)
|
||
|
|
call = ModelCall(
|
||
|
|
call_id=len(self.calls) + 1,
|
||
|
|
depth=depth,
|
||
|
|
role=role,
|
||
|
|
problem=problem,
|
||
|
|
input_text=input_text,
|
||
|
|
output_text=output_text,
|
||
|
|
input_tokens=token_count(self.tokenizer, input_text + forced_output_prefix),
|
||
|
|
output_tokens=token_count(self.tokenizer, generated_output),
|
||
|
|
forced_output_prefix=forced_output_prefix,
|
||
|
|
)
|
||
|
|
self.calls.append(call)
|
||
|
|
own_tokens = call.input_tokens + call.output_tokens
|
||
|
|
|
||
|
|
try:
|
||
|
|
action_type, paths = parse_model_output(output_text)
|
||
|
|
validate_action_protocol(output_text, action_type, role)
|
||
|
|
call.action_type = action_type
|
||
|
|
except ValueError as exc:
|
||
|
|
call.parse_error = str(exc)
|
||
|
|
return NodeResult(
|
||
|
|
final_answer=None,
|
||
|
|
trace=output_text,
|
||
|
|
critical_path_tokens=own_tokens,
|
||
|
|
parse_errors=[f"call_{call.call_id}: {exc}"],
|
||
|
|
stopped_reason="parse_error",
|
||
|
|
)
|
||
|
|
|
||
|
|
if action_type in {"SOLVE_DIRECTLY", "AGGREGATE"}:
|
||
|
|
return NodeResult(
|
||
|
|
final_answer=extract_final_answer(output_text),
|
||
|
|
trace=output_text,
|
||
|
|
critical_path_tokens=own_tokens,
|
||
|
|
parse_errors=[],
|
||
|
|
)
|
||
|
|
|
||
|
|
child_results = [
|
||
|
|
self._solve_node(problem=path, depth=depth + 1, role=f"path_{idx}")
|
||
|
|
for idx, path in enumerate(paths, start=1)
|
||
|
|
]
|
||
|
|
aggregate_problem = build_aggregation_problem(problem, paths, child_results)
|
||
|
|
aggregate_result = self._solve_node(
|
||
|
|
problem=aggregate_problem,
|
||
|
|
depth=depth + 1,
|
||
|
|
role="aggregate",
|
||
|
|
)
|
||
|
|
child_trace = "\n\n".join(
|
||
|
|
f"<problem>\n{path.strip()}\n</problem>\n{result.trace.strip()}"
|
||
|
|
for path, result in zip(paths, child_results)
|
||
|
|
)
|
||
|
|
full_trace = "\n\n".join(
|
||
|
|
part
|
||
|
|
for part in [
|
||
|
|
output_text.strip(),
|
||
|
|
child_trace.strip(),
|
||
|
|
f"<problem>\n{aggregate_problem.strip()}\n</problem>\n{aggregate_result.trace.strip()}",
|
||
|
|
]
|
||
|
|
if part
|
||
|
|
)
|
||
|
|
child_critical = max((result.critical_path_tokens for result in child_results), default=0)
|
||
|
|
parse_errors: list[str] = []
|
||
|
|
stopped_reason = "completed"
|
||
|
|
for result in child_results + [aggregate_result]:
|
||
|
|
parse_errors.extend(result.parse_errors)
|
||
|
|
if result.stopped_reason != "completed":
|
||
|
|
stopped_reason = result.stopped_reason
|
||
|
|
|
||
|
|
return NodeResult(
|
||
|
|
final_answer=aggregate_result.final_answer,
|
||
|
|
trace=full_trace,
|
||
|
|
critical_path_tokens=own_tokens + child_critical + aggregate_result.critical_path_tokens,
|
||
|
|
parse_errors=parse_errors,
|
||
|
|
stopped_reason=stopped_reason,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load_model_and_tokenizer(model_path: str, dtype: str = "bf16", device_map: str | None = None) -> tuple[Any, Any]:
|
||
|
|
import torch
|
||
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
||
|
|
if tokenizer.pad_token is None:
|
||
|
|
tokenizer.pad_token = tokenizer.eos_token
|
||
|
|
|
||
|
|
dtype_map = {
|
||
|
|
"bf16": torch.bfloat16,
|
||
|
|
"fp16": torch.float16,
|
||
|
|
"fp32": torch.float32,
|
||
|
|
"auto": "auto",
|
||
|
|
}
|
||
|
|
torch_dtype = dtype_map.get(dtype, torch.bfloat16)
|
||
|
|
kwargs: dict[str, Any] = {
|
||
|
|
"torch_dtype": torch_dtype,
|
||
|
|
"trust_remote_code": True,
|
||
|
|
}
|
||
|
|
if device_map:
|
||
|
|
kwargs["device_map"] = device_map
|
||
|
|
model = AutoModelForCausalLM.from_pretrained(model_path, **kwargs)
|
||
|
|
if not device_map:
|
||
|
|
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
|
||
|
|
# SFT checkpoints commonly persist use_cache=False from gradient checkpointing.
|
||
|
|
# Generation should always use the KV cache, or long rollouts become prohibitively slow.
|
||
|
|
model.config.use_cache = True
|
||
|
|
model.eval()
|
||
|
|
return model, tokenizer
|
||
|
|
|
||
|
|
|
||
|
|
def row_to_eval_record(
|
||
|
|
item: dict[str, Any],
|
||
|
|
*,
|
||
|
|
idx: int,
|
||
|
|
split: str,
|
||
|
|
problem_field: str | None,
|
||
|
|
answer_field: str | None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
p_field = problem_field or infer_field(item, ["problem", "question", "input"])
|
||
|
|
a_field = answer_field or infer_field(item, ["answer", "target", "final_answer", "solution"])
|
||
|
|
if p_field is None:
|
||
|
|
raise ValueError(f"Could not infer problem field. Available fields: {list(item.keys())}")
|
||
|
|
return {
|
||
|
|
"id": str(item.get("id", item.get("unique_id", f"{split}_{idx:06d}"))),
|
||
|
|
"problem": str(item[p_field]),
|
||
|
|
"answer": str(item[a_field]) if a_field and item.get(a_field) is not None else None,
|
||
|
|
"raw": dict(item),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def load_eval_jsonl(
|
||
|
|
*,
|
||
|
|
path: str | Path,
|
||
|
|
problem_field: str | None,
|
||
|
|
answer_field: str | None,
|
||
|
|
limit: int | None,
|
||
|
|
start_index: int = 0,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
rows: list[dict[str, Any]] = []
|
||
|
|
path = Path(path)
|
||
|
|
with path.open("r", encoding="utf-8") as handle:
|
||
|
|
for absolute_idx, line in enumerate(handle):
|
||
|
|
if absolute_idx < start_index:
|
||
|
|
continue
|
||
|
|
if limit is not None and len(rows) >= limit:
|
||
|
|
break
|
||
|
|
if not line.strip():
|
||
|
|
continue
|
||
|
|
item = json.loads(line)
|
||
|
|
rows.append(
|
||
|
|
row_to_eval_record(
|
||
|
|
item,
|
||
|
|
idx=absolute_idx,
|
||
|
|
split="jsonl",
|
||
|
|
problem_field=problem_field,
|
||
|
|
answer_field=answer_field,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def load_eval_dataset(
|
||
|
|
*,
|
||
|
|
dataset: str,
|
||
|
|
dataset_config: str | None,
|
||
|
|
split: str,
|
||
|
|
problem_field: str | None,
|
||
|
|
answer_field: str | None,
|
||
|
|
limit: int | None,
|
||
|
|
start_index: int = 0,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
ds = load_dataset(dataset, dataset_config, split=split) if dataset_config else load_dataset(dataset, split=split)
|
||
|
|
if start_index:
|
||
|
|
ds = ds.select(range(start_index, len(ds)))
|
||
|
|
if limit is not None:
|
||
|
|
ds = ds.select(range(min(limit, len(ds))))
|
||
|
|
|
||
|
|
rows: list[dict[str, Any]] = []
|
||
|
|
for idx, item in enumerate(ds, start=start_index):
|
||
|
|
rows.append(
|
||
|
|
row_to_eval_record(
|
||
|
|
dict(item),
|
||
|
|
idx=idx,
|
||
|
|
split=split,
|
||
|
|
problem_field=problem_field,
|
||
|
|
answer_field=answer_field,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def infer_field(item: dict[str, Any], candidates: list[str]) -> str | None:
|
||
|
|
for candidate in candidates:
|
||
|
|
if candidate in item:
|
||
|
|
return candidate
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def write_jsonl(path: str | Path, rows: list[dict[str, Any]]) -> None:
|
||
|
|
path = Path(path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with path.open("w", encoding="utf-8") as handle:
|
||
|
|
for row in rows:
|
||
|
|
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
|