mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
5d7fbac0ad
* math prompt improvements * ignore brackets in complex_arithmetic results * improve additional instruction in prompt of polynomial_equations * more strict tests for score_answer in polynomial_equations * simplify special reward handling * fix test_intermediate_integration * fix sokoban dataset * add common dataset score_answer consistency test
117 lines
4.8 KiB
Python
117 lines
4.8 KiB
Python
"""Sentence re-ordering task generator"""
|
|
|
|
import re
|
|
from dataclasses import dataclass
|
|
from random import Random
|
|
from typing import Any, Optional
|
|
|
|
from ..data import read_data_file
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
|
|
|
|
@dataclass
|
|
class SentenceReorderingConfig:
|
|
"""Configuration for sentence reordering task generation"""
|
|
|
|
min_words_in_sentence: int = 3
|
|
max_words_in_sentence: int = 20
|
|
seed: Optional[int] = None
|
|
size: int = 500
|
|
|
|
def validate(self) -> None:
|
|
"""Validate configuration parameters"""
|
|
assert self.min_words_in_sentence > 0, "min_words_in_sentence must be positive"
|
|
assert (
|
|
self.max_words_in_sentence >= self.min_words_in_sentence
|
|
), "max_words_in_sentence must be >= min_words_in_sentence"
|
|
assert (
|
|
self.max_words_in_sentence >= self.min_words_in_sentence
|
|
), "max_words_in_sentence must be >= min_words_in_sentence"
|
|
|
|
|
|
class SentenceReorderingDataset(ProceduralDataset):
|
|
"""Generates sentence reordering tasks from text spans"""
|
|
|
|
def __init__(self, config: SentenceReorderingConfig):
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
# Load and preprocess text
|
|
text = read_data_file("in_the_year_2889.txt")
|
|
# Extract sentences make sure they are greater than or equal to the number of words in a sentence
|
|
# Ensure that only the length of alphanumeric characters in the sentence is considered
|
|
self.sentences = [
|
|
sentence.strip()
|
|
for sentence in re.findall(r"[^.!?]+[.!?]", text) # Changed pattern to include the ending punctuation
|
|
if self.config.min_words_in_sentence
|
|
<= len(re.findall(r"\b\w+\b", sentence))
|
|
<= self.config.max_words_in_sentence
|
|
]
|
|
|
|
def _generate_sentence_dataset(self, sentence: str, seed: int, idx: int, shuffle=True):
|
|
"""
|
|
Generate a procedural dataset by shuffling the words in the input sentence.
|
|
|
|
Args:
|
|
sentence (str): The correct sentence to use for dataset generation.
|
|
seed (int): The seed to use for random number generation.
|
|
idx (int): The index to add to the seed for random number generation.
|
|
shuffle (bool): Whether to shuffle the words to create the input sentence.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the input sentence and the correct sentence (goal).
|
|
"""
|
|
rng = Random(seed + idx)
|
|
words = sentence.split() # Split the sentence into words
|
|
scrambled_words = words.copy()
|
|
if shuffle:
|
|
rng.shuffle(scrambled_words) # Shuffle the words to generate the input
|
|
input_sentence = " ".join(scrambled_words)
|
|
goal_sentence = " ".join(words)
|
|
return {"input": input_sentence, "goal": goal_sentence}
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
"""Generate a single sentence reordering task"""
|
|
rng = Random(self.seed + idx)
|
|
sentence_dataset = self._generate_sentence_dataset(rng.choice(self.sentences), self.seed, idx)
|
|
|
|
# Ensure only 'input' and 'goal' keys are present
|
|
if set(sentence_dataset.keys()) != {"input", "goal"}:
|
|
raise KeyError("The dictionary must contain only 'input' and 'goal' keys")
|
|
|
|
# Solve the task by sorting words to match the goal sentence
|
|
input_words = sentence_dataset["input"].split()
|
|
question = " ".join(input_words)
|
|
goal_words = sentence_dataset["goal"].split()
|
|
solved_sentence = " ".join(sorted(input_words, key=lambda word: goal_words.index(word)))
|
|
# Check for length of alphanumeric characters in the solved sentence
|
|
word_count = len(re.findall(r"\b\w+\b", solved_sentence))
|
|
|
|
return {
|
|
"question": f"Restore the correct order of words in the following sentence: {question}",
|
|
"answer": solved_sentence,
|
|
"metadata": {"word_count": word_count},
|
|
}
|
|
|
|
def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
|
|
reward = 0.0
|
|
expected_answer = entry["answer"]
|
|
if answer is not None:
|
|
try:
|
|
if expected_answer == answer:
|
|
return 1.0
|
|
goal_words = expected_answer.split()
|
|
answer_words = answer.split()
|
|
if len(goal_words) == len(answer_words):
|
|
credit = [
|
|
1 if goal_word.lower() == answer_word.lower() else 0
|
|
for goal_word, answer_word in zip(goal_words, answer_words)
|
|
]
|
|
reward = sum(credit) / len(credit)
|
|
else:
|
|
reward = 0.05
|
|
except:
|
|
reward = 0.0
|
|
return reward
|
|
|
|
|
|
register_dataset("sentence_reordering", SentenceReorderingDataset, SentenceReorderingConfig)
|