313 lines
12 KiB
Python
313 lines
12 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_PLAN_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"setting_title": {"type": "string"},
|
|
"setting_premise": {"type": "string"},
|
|
"style_guide": {"type": "string"},
|
|
},
|
|
"required": ["setting_title", "setting_premise", "style_guide"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
RENDER_CHUNK_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {"text": {"type": "string"}},
|
|
"required": ["text"],
|
|
"additionalProperties": False,
|
|
}
|
|
|
|
|
|
def _surface_names(skeleton: WorldSkeleton) -> dict[str, str]:
|
|
first = ("Zeru", "Parn", "Moki", "Tavi", "Luma", "Neri", "Kado", "Vela", "Sori", "Denu")
|
|
second = ("an", "or", "is", "um", "el", "ai", "on", "ar", "en", "ul")
|
|
names: dict[str, str] = {}
|
|
counters: dict[str, int] = {}
|
|
for entity in skeleton.entities:
|
|
index = counters.get(entity.type, 0)
|
|
counters[entity.type] = index + 1
|
|
if entity.type == "Entity":
|
|
names[entity.id] = first[index % len(first)] + second[(index // len(first)) % len(second)]
|
|
elif entity.type == "Set":
|
|
names[entity.id] = f"Collection-{index + 1:02d}"
|
|
elif entity.type == "Region":
|
|
names[entity.id] = f"Chamber-{index + 1:02d}"
|
|
else:
|
|
names[entity.id] = f"Gauge-{index + 1:02d}"
|
|
return names
|
|
|
|
|
|
def render_skeleton(
|
|
client: JsonModel,
|
|
skeleton: WorldSkeleton,
|
|
*,
|
|
render_chunk_size: int = 60,
|
|
) -> str:
|
|
if render_chunk_size < 1:
|
|
raise ValueError("render_chunk_size must be positive")
|
|
payload = skeleton.renderer_payload()
|
|
names = _surface_names(skeleton)
|
|
type_counts: dict[str, int] = {}
|
|
for entity_type in payload["entity_types"].values():
|
|
type_counts[entity_type] = type_counts.get(entity_type, 0) + 1
|
|
plan = client.request_json(
|
|
system=(
|
|
"Design a concise fictional setting for a large symbolic event history. Return only a title, a short "
|
|
"premise, and a style guide that maps the supplied relation and event vocabulary to consistent natural "
|
|
"language. Do not invent world facts, rules, probabilities, a query answer, or a proof."
|
|
),
|
|
user=json.dumps(
|
|
{
|
|
"entity_type_counts": type_counts,
|
|
"relations": [
|
|
"LOCATED_IN",
|
|
"ABOVE",
|
|
"MEMBER",
|
|
"SUBSET",
|
|
"DISJOINT",
|
|
"CONNECTED",
|
|
],
|
|
"event_kinds": ["MOVE", "UPDATE_NUMBER", "SET_STATE"],
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
schema_name="rendering_plan",
|
|
schema=RENDER_PLAN_SCHEMA,
|
|
)
|
|
if not all(isinstance(plan.get(key), str) and plan[key].strip() for key in RENDER_PLAN_SCHEMA["required"]):
|
|
raise RuntimeError("renderer returned an incomplete rendering plan")
|
|
|
|
records = [
|
|
{"section": "initial_observation", **item}
|
|
for item in skeleton.initial_state["public_observations"]
|
|
]
|
|
records.extend({"section": "event", **asdict(event)} for event in skeleton.events)
|
|
rendered_chunks: list[str] = []
|
|
for start in range(0, len(records), render_chunk_size):
|
|
chunk = records[start : start + render_chunk_size]
|
|
response = client.request_json(
|
|
system=(
|
|
"Render the supplied slice of one world history as compact natural-language log lines. Use the "
|
|
"provided setting, style guide, and exact surface-name map. Render every record exactly once and "
|
|
"begin its line with the record's I#### or V##### audit ID. Preserve all participants, relation "
|
|
"directions, values, before/after states, and event order. Do not summarize, infer rules, answer a "
|
|
"query, or add facts. Return only the finished lines for this slice."
|
|
),
|
|
user=json.dumps(
|
|
{
|
|
"rendering_plan": plan,
|
|
"surface_names": names,
|
|
"records": chunk,
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
schema_name="rendered_history_chunk",
|
|
schema=RENDER_CHUNK_SCHEMA,
|
|
)
|
|
text = response.get("text")
|
|
if not isinstance(text, str) or not text.strip():
|
|
raise RuntimeError(f"renderer returned an empty history chunk beginning at record {start}")
|
|
rendered_chunks.append(text.strip())
|
|
|
|
query_names = [names.get(argument, argument) for argument in skeleton.query.arguments]
|
|
if skeleton.query.predicate == "ABOVE":
|
|
question = f"At the end of the history, is {query_names[0]} above {query_names[1]}?"
|
|
elif skeleton.query.predicate == "MEMBER":
|
|
question = f"At the end of the history, is {query_names[0]} a member of {query_names[1]}?"
|
|
else:
|
|
raise RuntimeError(f"unsupported rendered query predicate: {skeleton.query.predicate}")
|
|
problem = (
|
|
f"# {plan['setting_title'].strip()}\n\n{plan['setting_premise'].strip()}\n\n"
|
|
+ "\n".join(rendered_chunks)
|
|
+ f"\n\nQuestion: {question}\nAnswer labels: {', '.join(skeleton.answer_options)}"
|
|
)
|
|
expected_ids = [item["id"] for item in records]
|
|
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,
|
|
render_chunk_size: int = 60,
|
|
) -> dict[str, Any]:
|
|
if render_chunk_size < 1:
|
|
raise ValueError("render_chunk_size must be positive")
|
|
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,
|
|
render_chunk_size=render_chunk_size,
|
|
)
|
|
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,
|
|
"render_chunk_size": render_chunk_size,
|
|
"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
|