mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-19 12:58:07 +00:00
81 lines
3 KiB
Python
81 lines
3 KiB
Python
"""Caesar cipher task generator"""
|
|
|
|
from dataclasses import dataclass
|
|
from random import Random
|
|
from typing import Optional
|
|
|
|
from ..data import read_data_file
|
|
from ..factory import ProceduralDataset, register_dataset
|
|
|
|
|
|
@dataclass
|
|
class CaesarCipherConfig:
|
|
"""Configuration for Caesar cipher task generation"""
|
|
|
|
delimiter: str = "." # Delimiter for splitting text into sentences
|
|
min_words: int = 3 # Minimum words per sentence
|
|
max_words: int = 20 # Maximum words per sentence
|
|
min_rotation: int = 1 # Minimum Caesar rotation
|
|
max_rotation: int = 25 # Maximum Caesar rotation
|
|
seed: Optional[int] = None
|
|
size: int = 500 # Virtual dataset size
|
|
|
|
def validate(self) -> None:
|
|
"""Validate configuration parameters"""
|
|
assert self.min_words > 0, "min_words must be positive"
|
|
assert self.max_words >= self.min_words, "max_words must be >= min_words"
|
|
assert 0 < self.min_rotation <= self.max_rotation < 26, "rotation must be in range [1,25]"
|
|
|
|
|
|
class CaesarCipherDataset(ProceduralDataset):
|
|
"""Generates Caesar cipher encryption/decryption tasks"""
|
|
|
|
def __init__(self, config: CaesarCipherConfig):
|
|
super().__init__(config=config, seed=config.seed, size=config.size)
|
|
|
|
# Load and preprocess text
|
|
text = read_data_file("in_the_year_2889.txt")
|
|
|
|
# Split into sentences and filter
|
|
sentences = [s.strip() for s in text.split(config.delimiter) if s.strip()]
|
|
|
|
# Process each sentence
|
|
self.valid_sentences = []
|
|
for sentence in sentences:
|
|
# Split into words and filter for alpha-only
|
|
words = [w.upper() for w in sentence.split() if w.isalpha()]
|
|
if self.config.min_words <= len(words) <= self.config.max_words:
|
|
self.valid_sentences.append(" ".join(words))
|
|
|
|
def _caesar_encrypt(self, text: str, rotation: int) -> str:
|
|
"""Apply Caesar cipher encryption with given rotation"""
|
|
result = []
|
|
for char in text:
|
|
if char.isalpha():
|
|
# Convert to 0-25 range, rotate, convert back to ASCII
|
|
base = ord("A")
|
|
rotated = (ord(char) - base + rotation) % 26
|
|
result.append(chr(base + rotated))
|
|
else:
|
|
result.append(char)
|
|
return "".join(result)
|
|
|
|
def __getitem__(self, idx: int) -> dict:
|
|
"""Generate a single Caesar cipher task"""
|
|
rng = Random(self.seed + idx)
|
|
|
|
# Select random sentence and rotation
|
|
sentence = rng.choice(self.valid_sentences)
|
|
rotation = rng.randint(self.config.min_rotation, self.config.max_rotation)
|
|
|
|
# Generate cipher text
|
|
cipher_text = self._caesar_encrypt(sentence, rotation)
|
|
|
|
return {
|
|
"question": f"Decrypt this Caesar cipher text: {cipher_text}. Provide only the decrypted text as your final answer.",
|
|
"answer": sentence,
|
|
"metadata": {"rotation": rotation, "cipher_text": cipher_text, "clear_text": sentence},
|
|
}
|
|
|
|
|
|
register_dataset("caesar_cipher", CaesarCipherDataset, CaesarCipherConfig)
|