mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-30 17:40:45 +00:00
test: Add comprehensive unit tests for RushHourDataset
This commit is contained in:
parent
f0ad47a29a
commit
ca5e23e195
2 changed files with 75 additions and 19 deletions
|
|
@ -1,19 +1,19 @@
|
|||
import random
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Optional, Tuple
|
||||
import random
|
||||
|
||||
from ..data import read_data_file
|
||||
from ..factory import ProceduralDataset, register_dataset
|
||||
|
||||
TEST_STRING = "BBoKMxDDDKMoIAALooIoJLEEooJFFNoGGoxN"
|
||||
|
||||
@dataclass
|
||||
class RushHourConfig:
|
||||
"""Configuration for Rush Hour puzzle generation"""
|
||||
|
||||
min_moves: int = 1
|
||||
max_moves: int = 50
|
||||
seed: Optional[int] = None
|
||||
seed: Optional[int] = None
|
||||
size: int = 500
|
||||
|
||||
def validate(self) -> None:
|
||||
|
|
@ -100,10 +100,10 @@ class RushHourDataset(ProceduralDataset):
|
|||
|
||||
def __init__(self, config: RushHourConfig):
|
||||
super().__init__(config=config, seed=config.seed, size=config.size)
|
||||
|
||||
|
||||
# Load and filter puzzles from data file
|
||||
self.puzzles: List[Tuple[str, int]] = [] # (board_config, min_moves)
|
||||
|
||||
|
||||
data = read_data_file("rush_18k.txt")
|
||||
for line in data.splitlines():
|
||||
if not line.strip():
|
||||
|
|
@ -112,21 +112,19 @@ class RushHourDataset(ProceduralDataset):
|
|||
if len(parts) >= 2:
|
||||
min_moves = int(parts[0])
|
||||
board_config = parts[1]
|
||||
|
||||
|
||||
if config.min_moves <= min_moves <= config.max_moves:
|
||||
self.puzzles.append((board_config, min_moves))
|
||||
|
||||
|
||||
if not self.puzzles:
|
||||
raise ValueError(
|
||||
f"No puzzles found with moves between {config.min_moves} and {config.max_moves}"
|
||||
)
|
||||
raise ValueError(f"No puzzles found with moves between {config.min_moves} and {config.max_moves}")
|
||||
|
||||
def __getitem__(self, idx: int) -> dict:
|
||||
"""Generate a single Rush Hour puzzle
|
||||
|
||||
|
||||
Args:
|
||||
idx: Index of the item to generate
|
||||
|
||||
|
||||
Returns:
|
||||
dict with keys:
|
||||
- question: str, the formatted board with instructions
|
||||
|
|
@ -135,23 +133,23 @@ class RushHourDataset(ProceduralDataset):
|
|||
"""
|
||||
# Create deterministic RNG from base seed and idx
|
||||
rng = random.Random(self.seed + idx)
|
||||
|
||||
|
||||
# Randomly select a puzzle meeting our criteria
|
||||
board_config, min_moves = rng.choice(self.puzzles)
|
||||
|
||||
|
||||
# Create board and get string representation
|
||||
board = Board(board_config)
|
||||
board_display = board.board_str()
|
||||
|
||||
|
||||
instructions = (
|
||||
"Move the red car (AA) to the exit on the right.\n"
|
||||
"Specify moves in the format: 'F+1 K+1 M-1 C+3 H+2 ...'\n"
|
||||
"where the letter is the vehicle and +/- number is spaces to move right/left or down/up."
|
||||
)
|
||||
|
||||
|
||||
return {
|
||||
"question": f"{instructions}\n\nBoard:\n{board_display}",
|
||||
"answer": "", # Multiple valid solutions exist
|
||||
"answer": None, # Multiple valid solutions exist
|
||||
"metadata": {
|
||||
"board_config": board_config,
|
||||
"min_moves": min_moves,
|
||||
|
|
@ -330,4 +328,3 @@ class Board:
|
|||
|
||||
# Register the dataset
|
||||
register_dataset("rush_hour", RushHourDataset, RushHourConfig)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue