215 lines
8.4 KiB
Python
215 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Protocol
|
|
|
|
from .models import WorldSkeleton
|
|
from .skeletons import generate_skeletons
|
|
|
|
|
|
class JsonModel(Protocol):
|
|
def request_json(
|
|
self,
|
|
*,
|
|
system: str,
|
|
user: str,
|
|
schema_name: str,
|
|
schema: dict[str, Any],
|
|
) -> dict[str, Any]: ...
|
|
|
|
def request_text(self, *, system: str, user: str) -> str: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class EvaluationResult:
|
|
skeleton_id: str
|
|
family: str
|
|
complexity: int
|
|
rendered_problem: str
|
|
answer_options: tuple[str, ...]
|
|
gold_answer: str
|
|
solver_output: str
|
|
model_answer: str | None
|
|
format_compliant: bool
|
|
parse_error: str | None
|
|
correct: bool
|
|
|
|
|
|
RENDER_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {"problem": {"type": "string"}},
|
|
"required": ["problem"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
|
|
def render_skeleton(client: JsonModel, skeleton: WorldSkeleton) -> str:
|
|
prompt = json.dumps(skeleton.renderer_payload(), ensure_ascii=False, indent=2)
|
|
response = client.request_json(
|
|
system=(
|
|
"Turn a large formal history from one simulated world into a self-contained natural-language "
|
|
"reasoning problem. Invent one coherent concrete or fictional setting and assign stable names to "
|
|
"every symbol. Render every initial observation and every realized event exactly once, preserving "
|
|
"time order, participants, relation direction, numeric values, and state changes. Render the query "
|
|
"at the end. The events are facts that actually occurred, not repeated probability experiments. "
|
|
"Never mention the generator, hidden program, probabilities, proof path, gold answer, or authoring "
|
|
"instructions. Do not directly state the queried proposition. Begin each rendered fact or event line "
|
|
"with its original I#### or V##### audit ID. Use compact chronological log lines rather than "
|
|
"explanatory prose, but do not omit apparent distractors. State the answer labels; the "
|
|
"evaluator separately supplies the solver response protocol."
|
|
),
|
|
user=f"Render this intermediate representation:\n{prompt}",
|
|
schema_name="rendered_problem",
|
|
schema=RENDER_SCHEMA,
|
|
)
|
|
problem = response.get("problem")
|
|
if not isinstance(problem, str) or not problem.strip():
|
|
raise RuntimeError("renderer returned an empty problem")
|
|
problem = problem.strip()
|
|
expected_ids = [item["id"] for item in skeleton.initial_state["public_observations"]]
|
|
expected_ids.extend(event.id for event in skeleton.events)
|
|
missing = [identifier for identifier in expected_ids if identifier not in problem]
|
|
if missing:
|
|
preview = ", ".join(missing[:10])
|
|
raise RuntimeError(
|
|
f"renderer omitted {len(missing)} initial observations or events; first missing IDs: {preview}"
|
|
)
|
|
return problem
|
|
|
|
|
|
def parse_marked_answer(output: str, answer_options: tuple[str, ...]) -> str:
|
|
lines = [line.strip() for line in output.splitlines() if line.strip()]
|
|
if not lines or not lines[-1].startswith("####"):
|
|
raise RuntimeError("solver output must end with '#### <answer>'")
|
|
answer = lines[-1][4:].strip()
|
|
if answer not in answer_options:
|
|
raise RuntimeError(f"solver returned invalid answer label: {answer!r}")
|
|
return answer
|
|
|
|
|
|
def parse_solver_answer(output: str, answer_options: tuple[str, ...]) -> tuple[str, bool]:
|
|
"""Prefer the marker protocol, then recover one unambiguous label from the final line."""
|
|
try:
|
|
return parse_marked_answer(output, answer_options), True
|
|
except RuntimeError as strict_error:
|
|
lines = [line.strip() for line in output.splitlines() if line.strip()]
|
|
if not lines:
|
|
raise strict_error
|
|
final_line = lines[-1]
|
|
matches = [
|
|
option
|
|
for option in answer_options
|
|
if re.search(rf"(?<![A-Za-z0-9_]){re.escape(option)}(?![A-Za-z0-9_])", final_line)
|
|
]
|
|
if len(matches) == 1:
|
|
return matches[0], False
|
|
raise strict_error
|
|
|
|
|
|
def solve_problem(client: JsonModel, problem: str, answer_options: tuple[str, ...]) -> tuple[str, str]:
|
|
output = request_solver_output(client, problem, answer_options)
|
|
answer, _format_compliant = parse_solver_answer(output, answer_options)
|
|
return output, answer
|
|
|
|
|
|
def request_solver_output(client: JsonModel, problem: str, answer_options: tuple[str, ...]) -> str:
|
|
return client.request_text(
|
|
system=(
|
|
"Solve the question from the described world history. Track entity identity, relation direction, "
|
|
"temporal updates, set inheritance, ordering, and relevant compositions while ignoring distractors. "
|
|
"Use only the stated observations and ordinary meanings explicitly established by the story. Keep "
|
|
"the explanation under 200 words, then put "
|
|
"the final answer on the last non-empty line in the exact form '#### <answer>', where <answer> "
|
|
"is exactly one permitted answer label. Do not use unstated real-world assumptions."
|
|
),
|
|
user=f"{problem}\n\nPermitted answer labels: {', '.join(answer_options)}",
|
|
)
|
|
|
|
|
|
def _write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for record in records:
|
|
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
|
|
|
|
|
|
def run_evaluation(
|
|
*,
|
|
renderer: JsonModel,
|
|
solver: JsonModel,
|
|
count: int,
|
|
length: str | int,
|
|
seed: int,
|
|
output_dir: Path,
|
|
entity_count: int = 128,
|
|
simulation_steps: int = 1000,
|
|
) -> dict[str, Any]:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
skeletons = generate_skeletons(
|
|
count=count,
|
|
length=length,
|
|
seed=seed,
|
|
entity_count=entity_count,
|
|
simulation_steps=simulation_steps,
|
|
)
|
|
_write_jsonl(output_dir / "skeletons.jsonl", [item.to_dict() for item in skeletons])
|
|
|
|
results: list[EvaluationResult] = []
|
|
rendered_records: list[dict[str, str]] = []
|
|
solver_records: list[dict[str, str]] = []
|
|
for skeleton in skeletons:
|
|
problem = render_skeleton(renderer, skeleton)
|
|
rendered_records.append({"skeleton_id": skeleton.id, "rendered_problem": problem})
|
|
_write_jsonl(output_dir / "rendered.jsonl", rendered_records)
|
|
solver_output = request_solver_output(solver, problem, skeleton.answer_options)
|
|
solver_records.append({"skeleton_id": skeleton.id, "solver_output": solver_output})
|
|
_write_jsonl(output_dir / "solver-outputs.jsonl", solver_records)
|
|
try:
|
|
answer, format_compliant = parse_solver_answer(solver_output, skeleton.answer_options)
|
|
parse_error = None
|
|
except RuntimeError as exc:
|
|
answer = None
|
|
format_compliant = False
|
|
parse_error = str(exc)
|
|
results.append(
|
|
EvaluationResult(
|
|
skeleton_id=skeleton.id,
|
|
family=skeleton.family,
|
|
complexity=skeleton.complexity,
|
|
rendered_problem=problem,
|
|
answer_options=skeleton.answer_options,
|
|
gold_answer=skeleton.gold_answer,
|
|
solver_output=solver_output,
|
|
model_answer=answer,
|
|
format_compliant=format_compliant,
|
|
parse_error=parse_error,
|
|
correct=answer == skeleton.gold_answer,
|
|
)
|
|
)
|
|
_write_jsonl(output_dir / "results.jsonl", [asdict(item) for item in results])
|
|
|
|
correct = sum(item.correct for item in results)
|
|
parsed = sum(item.model_answer is not None for item in results)
|
|
format_compliant = sum(item.format_compliant for item in results)
|
|
summary = {
|
|
"count": len(results),
|
|
"correct": correct,
|
|
"accuracy": correct / len(results),
|
|
"parsed": parsed,
|
|
"parse_rate": parsed / len(results),
|
|
"format_compliant": format_compliant,
|
|
"format_compliance": format_compliant / len(results),
|
|
"length": length,
|
|
"complexity": skeletons[0].complexity,
|
|
"seed": seed,
|
|
"entity_count": entity_count,
|
|
"simulation_steps": simulation_steps,
|
|
"event_counts": [len(item.events) for item in skeletons],
|
|
}
|
|
(output_dir / "summary.json").write_text(
|
|
json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
|
|
)
|
|
return summary
|