112 lines
4.0 KiB
Python
112 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from collections.abc import Callable
|
|
|
|
from .models import Fact, Query, Skeleton
|
|
|
|
|
|
LENGTH_PRESETS = {"short": 2, "medium": 4, "long": 6}
|
|
|
|
|
|
def resolve_length(value: str | int) -> int:
|
|
"""Resolve a named preset or an explicit number of proof-bearing operations."""
|
|
if isinstance(value, int):
|
|
result = value
|
|
elif value in LENGTH_PRESETS:
|
|
result = LENGTH_PRESETS[value]
|
|
else:
|
|
try:
|
|
result = int(value)
|
|
except ValueError as exc:
|
|
choices = ", ".join(LENGTH_PRESETS)
|
|
raise ValueError(f"length must be an integer or one of: {choices}") from exc
|
|
if not 1 <= result <= 20:
|
|
raise ValueError("length must be between 1 and 20")
|
|
return result
|
|
|
|
|
|
def _order_chain(index: int, complexity: int, rng: random.Random) -> Skeleton:
|
|
nodes = [f"E{i}" for i in range(complexity + 1)]
|
|
facts = [Fact("LEFT_OF", (nodes[i], nodes[i + 1])) for i in range(complexity)]
|
|
rng.shuffle(facts)
|
|
return Skeleton(
|
|
id=f"order-{index:04d}",
|
|
family="transitive_order",
|
|
complexity=complexity,
|
|
entities={node: "Entity" for node in nodes},
|
|
facts=tuple(facts),
|
|
rules=("LEFT_OF is transitive", "RIGHT_OF(x,y) iff LEFT_OF(y,x)"),
|
|
query=Query("RELATION", (nodes[0], nodes[-1])),
|
|
answer_options=("LEFT_OF", "RIGHT_OF", "UNKNOWN"),
|
|
gold_answer="LEFT_OF",
|
|
proof=tuple(f"LEFT_OF({nodes[i]},{nodes[i + 1]})" for i in range(complexity)),
|
|
)
|
|
|
|
|
|
def _set_chain(index: int, complexity: int, rng: random.Random) -> Skeleton:
|
|
item = "E0"
|
|
sets = [f"S{i}" for i in range(complexity)]
|
|
facts = [Fact("MEMBER_OF", (item, sets[0]))]
|
|
facts.extend(Fact("SUBSET_OF", (sets[i], sets[i + 1])) for i in range(complexity - 1))
|
|
rng.shuffle(facts)
|
|
return Skeleton(
|
|
id=f"set-{index:04d}",
|
|
family="membership_inheritance",
|
|
complexity=complexity,
|
|
entities={item: "Entity", **{name: "Set" for name in sets}},
|
|
facts=tuple(facts),
|
|
rules=("MEMBER_OF(x,A) and SUBSET_OF(A,B) imply MEMBER_OF(x,B)",),
|
|
query=Query("MEMBER_OF", (item, sets[-1])),
|
|
answer_options=("YES", "NO", "UNKNOWN"),
|
|
gold_answer="YES",
|
|
proof=tuple(
|
|
[f"MEMBER_OF({item},{sets[0]})"]
|
|
+ [f"SUBSET_OF({sets[i]},{sets[i + 1]})" for i in range(complexity - 1)]
|
|
),
|
|
)
|
|
|
|
|
|
def _state_updates(index: int, complexity: int, rng: random.Random) -> Skeleton:
|
|
start = rng.randint(3, 9)
|
|
operations: list[tuple[str, int]] = []
|
|
value = start
|
|
for _ in range(complexity):
|
|
amount = rng.randint(1, 4)
|
|
predicate = rng.choice(("INCREMENT", "DECREMENT"))
|
|
if predicate == "DECREMENT" and value - amount < 0:
|
|
predicate = "INCREMENT"
|
|
value = value + amount if predicate == "INCREMENT" else value - amount
|
|
operations.append((predicate, amount))
|
|
|
|
facts = [Fact("INITIAL_VALUE", ("E0", str(start)))]
|
|
facts.extend(Fact(predicate, ("E0", str(amount))) for predicate, amount in operations)
|
|
lower_distractor = value - 1 if value > 0 else value + 2
|
|
return Skeleton(
|
|
id=f"state-{index:04d}",
|
|
family="sequential_state_update",
|
|
complexity=complexity,
|
|
entities={"E0": "Counter"},
|
|
facts=tuple(facts),
|
|
rules=("Apply updates in the listed order",),
|
|
query=Query("FINAL_VALUE", ("E0",)),
|
|
answer_options=(str(value), str(value + 1), str(lower_distractor)),
|
|
gold_answer=str(value),
|
|
proof=tuple(f"{predicate}({amount})" for predicate, amount in operations),
|
|
)
|
|
|
|
|
|
GENERATORS: tuple[Callable[[int, int, random.Random], Skeleton], ...] = (
|
|
_order_chain,
|
|
_set_chain,
|
|
_state_updates,
|
|
)
|
|
|
|
|
|
def generate_skeletons(count: int, length: str | int, seed: int = 0) -> list[Skeleton]:
|
|
if count < 1:
|
|
raise ValueError("count must be positive")
|
|
complexity = resolve_length(length)
|
|
rng = random.Random(seed)
|
|
return [GENERATORS[i % len(GENERATORS)](i, complexity, rng) for i in range(count)]
|