reasoning-gym/reasoning_gym/algorithmic/string_insertion.py
Zafir Stojanovski dced3bfc45
fix(curriculum): Make boundaries in curriculum more sensible (#407)
* init

* fix tests

* unify codeio

* filtered for libraries not present in reasoning-gym

* fix more bounds

* puzzle24

* knight swap curriculum

* fix number sorting

* fix attributes

* add validation of config in creation of dataset

* dry run for instantiating and validating the datasets

* remove unused imports

* fix curriculum tests to reference newly updated attribute names
2025-04-04 20:24:14 +02:00

137 lines
5.1 KiB
Python

"""Insert into string according to a pattern
https://github.com/yongchao98/CodeSteer-v1.0/blob/main/create_dataset/create_dataset_string_insertion.py
"""
from dataclasses import dataclass
from random import Random
from typing import Any, Optional
from ..coaching import BaseCurriculum, RangeAttributeDefinition
from ..factory import ProceduralDataset, register_dataset
QUESTION_TEMPLATE = """Given a string consisting of characters A, B, C, D, and E, your job is to insert a character according to the following pattern:
1. If there is a substring ABCD in the string, insert the character A after the substring.
2. If there is a substring BCDE in the string, insert the character B after the substring.
3. If there is a substring CDEA in the string, insert the character C after the substring.
4. If there is a substring DEAB in the string, insert the character D after the substring.
5. If there is a substring EABC in the string, insert the character E after the substring.
Once you have inserted a character, you have to skip over the substring and the inserted character and continue the search from the next character.
Your output should be a string that has been modified according to the pattern.
Given the following string, provide the answer after inserting the characters according to the pattern: {string}
"""
DATASET_NAME = "string_insertion"
@dataclass
class StringInsertionConfig:
"""Configuration for String Insertion dataset generation"""
min_string_length: int = 5 # Minimum string length
max_string_length: int = 20 # Maximum string length
size: int = 500 # Virtual dataset size
seed: Optional[int] = None
def validate(self):
"""Validate configuration parameters"""
assert 5 <= self.min_string_length, "Minimum string length should be at least 5"
assert self.min_string_length <= self.max_string_length, "Minimum string length should be less than maximum"
class StringInsertionDataset(ProceduralDataset):
"""Generates String Insertion exercises with configurable difficulty"""
def __init__(self, config: StringInsertionConfig):
super().__init__(config=config, seed=config.seed, size=config.size)
self.vocabulary = ["A", "B", "C", "D", "E"]
self.insertion_rules = [
("ABCD", "A"),
("BCDE", "B"),
("CDEA", "C"),
("DEAB", "D"),
("EABC", "E"),
]
def _get_answer(self, string: str) -> str:
"""Apply insertion rules to a string"""
output = []
i = 0
while i < len(string):
inserted = False
for pattern, char in self.insertion_rules:
substring = string[i : i + len(pattern)]
if substring == pattern:
output.append(substring + char)
i += len(pattern)
inserted = True
break
if not inserted:
output.append(string[i])
i += 1
return "".join(output)
def score_answer(self, answer: Optional[str], entry: dict[str, Any]) -> float:
"""Overwrite this method in derived classes if a single oracle answer is not available."""
oracle_answer = entry["answer"]
if isinstance(answer, str):
if answer == oracle_answer:
return 1.0
else:
try:
# check if answer is python list of characters
answer = "".join(eval(answer))
if answer == oracle_answer:
return 0.1
except Exception:
pass
return 0.0
def __getitem__(self, idx: int) -> dict:
"""Generate a single String Insertion question"""
rng = Random(self.seed + idx)
string_length = rng.randint(self.config.min_string_length, self.config.max_string_length)
string = "".join(rng.choice(self.vocabulary) for _ in range(string_length))
answer = self._get_answer(string)
return {
"question": QUESTION_TEMPLATE.format(string=string),
"answer": str(answer),
"metadata": {
"source_dataset": DATASET_NAME,
"source_index": idx,
"string": string,
"solution": answer,
"string_length": string_length,
"difficulty": {
"string_length": (self.config.min_string_length, self.config.max_string_length),
},
},
}
class StringInsertionCurriculum(BaseCurriculum):
def __init__(self):
super().__init__(StringInsertionCurriculum.__name__, StringInsertionConfig)
# Define attributes
self._define_attributes(
RangeAttributeDefinition(
name="string_length",
levels=[10, 50, 100, 500],
description="Length of the string",
lower_field_name="min_string_length",
upper_field_name="max_string_length",
ensure_interval=True,
),
)
register_dataset(DATASET_NAME, StringInsertionDataset, StringInsertionConfig, StringInsertionCurriculum)