mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
from dataclasses import dataclass
|
|
from random import Random, seed
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
from .contrib.logic_puzzle.generate import generate_puzzle
|
|
|
|
|
|
@dataclass
|
|
class ZebraConfig:
|
|
"""Configuration for zebra puzzle generation"""
|
|
|
|
num_people: int = 4
|
|
num_characteristics: int = 4
|
|
seed: Optional[int] = None
|
|
size: int = 500
|
|
|
|
def validate(self):
|
|
"""Validate configuration parameters"""
|
|
assert 2 <= self.num_people <= 7, "num_people must be between 2 and 7"
|
|
assert 2 <= self.num_characteristics <= 7, "num_characteristics must be between 2 and 7"
|
|
|
|
|
|
class ZebraDataset(ProceduralDataset):
|
|
"""Generates [Zebra Puzzles](https://en.wikipedia.org/wiki/Zebra_Puzzle) with configurable parameters"""
|
|
|
|
def __init__(self, config: ZebraConfig):
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
"""Generate a single Zebra task
|
|
|
|
Returns:
|
|
dict with keys:
|
|
- question: str, the task description
|
|
- answer: str, a solution string
|
|
- metadata: dict with generation parameters
|
|
"""
|
|
seed(self.seed + idx)
|
|
|
|
K = self.config.num_people
|
|
M = self.config.num_characteristics
|
|
instance, puzzle = generate_puzzle(K, M, "train")
|
|
q = instance["questions"][0]["question"]
|
|
answer = instance["questions"][0]["answer"]
|
|
question = str(puzzle) + "\n" + q
|
|
|
|
return {
|
|
"question": question,
|
|
"answer": answer,
|
|
"metadata": {
|
|
"num_people": K,
|
|
"num_characteristics": M,
|
|
},
|
|
}
|
|
|
|
def score_answer(self, answer: Optional[str], entry: Dict[str, any]) -> float:
|
|
"""Determine if the solution provided solves the Zebra task.
|
|
|
|
The function awards 1.0 for a correct answer.
|
|
|
|
Args:
|
|
answer (Optional[str]): The user's answer.
|
|
entry (Dict[str, any]): The original dataset entry containing the correct answer.
|
|
|
|
Returns:
|
|
float: The computed score between 0.0 and 1.0.
|
|
"""
|
|
|
|
if answer == None:
|
|
return 0.0
|
|
if answer.lower().replace("\n", "") != entry["answer"].lower().replace("\n", ""):
|
|
return 0.01
|
|
else:
|
|
return 1.0 # Yay
|
|
|
|
|
|
register_dataset("zebra_puzzles", ZebraDataset, ZebraConfig)
|