mirror of
https://github.com/open-thought/reasoning-gym.git
synced 2026-04-22 16:49:06 +00:00
136 lines
5.2 KiB
Python
136 lines
5.2 KiB
Python
"""Tests for String Splitting questions generation"""
|
|
|
|
import pytest
|
|
|
|
from reasoning_gym.algorithmic.string_splitting import StringSplittingConfig, StringSplittingDataset
|
|
|
|
QUESTION_TEMPLATE = """There is a dismantling engineer who has old machines A, B, and C.
|
|
He discovered that he can obtain a batch of new parts X, Y, Z through the following rules:
|
|
1. One unit of machine A can be dismanteled into two units of part X and one unit of part Y.
|
|
2. Two units of machine B can be dismanteled into one unit of part X.
|
|
3. Two units of machine C can be dismanteled into one unit of part Y.
|
|
4. One unit of machine B and one unit of machine C can be combined into one unit of machine A.
|
|
5. One unit of part X and one unit of part Y can be combined into one unit of part Z.
|
|
|
|
Given a certain number of initial machines, your job is to continuously cycle through the rules 1-5 above, exausting one rule at a time, until no more rules can be applied, or until a state (counts of each machine and part type) is repeated.
|
|
After you make use of a rule, you should update the counts of each machine and part type accordingly, and then restart the process from rule 1.
|
|
|
|
The output should be the count of each machine and part type after the rules have been exhaustively applied in the following order: A B C X Y Z.
|
|
For example 1 0 1 5 4 3 means that you have 1 machine A, 0 machine B, 1 machine C, 5 part X, 4 part Y, and 3 part Z.
|
|
|
|
Example:
|
|
- Input: You have 2 machines A, 0 machines B, and 1 machine C.
|
|
- Output: 0 0 1 2 0 2
|
|
- Explanation
|
|
0. Initial state: 2 0 1 0 0 0
|
|
1. We can apply rule 1 and trade 1 machine A for 2 part X and 1 part Y: 1 0 1 2 1 0
|
|
2. Starting over, we can apply rule 1 again: 0 0 1 4 2 0
|
|
3. In the next iteration, we can apply rule 5 and trade 1 part X and 1 part Y for 1 part Z: 0 0 1 3 1 1
|
|
4. In the next iteration, we can apply rule 5 again: 0 0 1 2 0 2
|
|
5. We can't apply any more rules, so the final answer is 0 0 1 2 0 2
|
|
|
|
Now, you have {A_machine} machine A, {B_machine} machine B, and {C_machine} machine C. Provide the count of each machine and part type after applying the above rules.
|
|
"""
|
|
|
|
|
|
def test_string_splitting_config_validation():
|
|
"""Test that invalid configs raise appropriate errors"""
|
|
|
|
with pytest.raises(AssertionError):
|
|
config = StringSplittingConfig(min_initial_machines=-1) # negative not allowed
|
|
config.validate()
|
|
|
|
with pytest.raises(AssertionError):
|
|
config = StringSplittingConfig(min_initial_machines=3, max_initial_machines=2) # min > max
|
|
config.validate()
|
|
|
|
|
|
def test_string_splitting_dataset_deterministic():
|
|
"""Test that dataset generates same items with same seed"""
|
|
config = StringSplittingConfig(seed=42, size=10)
|
|
dataset1 = StringSplittingDataset(config)
|
|
dataset2 = StringSplittingDataset(config)
|
|
|
|
for i in range(len(dataset1)):
|
|
assert dataset1[i] == dataset2[i]
|
|
|
|
|
|
def test_string_splitting_dataset_items():
|
|
"""Test basic properties of generated items"""
|
|
config = StringSplittingConfig(min_initial_machines=1, max_initial_machines=5, size=10, seed=42)
|
|
dataset = StringSplittingDataset(config)
|
|
|
|
for i in range(len(dataset)):
|
|
item = dataset[i]
|
|
# Check item structure
|
|
assert isinstance(item, dict)
|
|
assert "question" in item
|
|
assert "answer" in item
|
|
assert "metadata" in item
|
|
|
|
# Check metadata
|
|
assert "states" in item["metadata"]
|
|
assert "solution" in item["metadata"]
|
|
|
|
states = item["metadata"]["states"]
|
|
solution = item["metadata"]["solution"]
|
|
|
|
# Verify dimensions
|
|
assert len(states) > 0
|
|
assert states[-1] == solution
|
|
for i in range(3):
|
|
assert 1 <= states[0][i] <= 5
|
|
for i in range(3, 6):
|
|
assert states[0][i] == 0
|
|
|
|
|
|
def test_string_splitting_dataset_iteration():
|
|
"""Test that iteration respects dataset size"""
|
|
config = StringSplittingConfig(size=5, seed=42)
|
|
dataset = StringSplittingDataset(config)
|
|
|
|
items = list(dataset)
|
|
assert len(items) == config.size
|
|
|
|
# Test multiple iterations yield same items
|
|
assert items == list(dataset)
|
|
|
|
|
|
def test_string_splitting_answer():
|
|
"""Test the answer calculation"""
|
|
config = StringSplittingConfig(seed=42)
|
|
dataset = StringSplittingDataset(config)
|
|
|
|
# Empty input
|
|
counts = [0, 0, 0, 0, 0, 0]
|
|
assert dataset._apply_rule(counts) == [0, 0, 0, 0, 0, 0]
|
|
|
|
# Rule 1: 1A -> 2X 1Y
|
|
counts = [1, 0, 0, 0, 0, 0]
|
|
assert dataset._apply_rule(counts) == [0, 0, 0, 2, 1, 0]
|
|
|
|
# Rule 2: 2B -> 1X
|
|
counts = [0, 2, 0, 0, 0, 0]
|
|
assert dataset._apply_rule(counts) == [0, 0, 0, 1, 0, 0]
|
|
|
|
# Rule 3: 2C -> 1Y
|
|
counts = [0, 0, 2, 0, 0, 0]
|
|
assert dataset._apply_rule(counts) == [0, 0, 0, 0, 1, 0]
|
|
|
|
# Rule 4: B + C -> A
|
|
counts = [0, 1, 1, 0, 0, 0]
|
|
assert dataset._apply_rule(counts) == [1, 0, 0, 0, 0, 0]
|
|
|
|
# Rule 5: X + Y -> Z
|
|
counts = [0, 0, 0, 1, 1, 0]
|
|
assert dataset._apply_rule(counts) == [0, 0, 0, 0, 0, 1]
|
|
|
|
# 1-shot example used in the prompt
|
|
A_machine, B_machine, C_machine = 2, 0, 1
|
|
assert dataset._get_answer(A_machine, B_machine, C_machine) == [
|
|
[2, 0, 1, 0, 0, 0],
|
|
[1, 0, 1, 2, 1, 0],
|
|
[0, 0, 1, 4, 2, 0],
|
|
[0, 0, 1, 3, 1, 1],
|
|
[0, 0, 1, 2, 0, 2],
|
|
]
|