gsm_symbolic generator changes

tests/test_gsm_symbolic.py

@@ -0,0 +1,92 @@
+from random import Random
+
+import pytest
+
+from reasoning_gym.arithmetic.gsm_symbolic import GSMSymbolicDataset, GSMSymbolicDatasetConfig
+
+
+def test_gsm_symbolic_config_validation():
+    """Test that config validation works"""
+    config = GSMSymbolicDatasetConfig(size=-1)
+    with pytest.raises(AssertionError):
+        config.validate()
+
+
+def test_gsm_symbolic_deterministic():
+    """Test that dataset generates same items with same seed"""
+    config = GSMSymbolicDatasetConfig(seed=42, size=10)
+    dataset1 = GSMSymbolicDataset(config)
+    dataset2 = GSMSymbolicDataset(config)
+
+    for i in range(len(dataset1)):
+        assert dataset1[i] == dataset2[i]
+
+
+def test_gsm_symbolic_items():
+    """Test basic properties of generated items"""
+    config = GSMSymbolicDatasetConfig(size=100, seed=42)
+    dataset = GSMSymbolicDataset(config)
+
+    for i in range(len(dataset)):
+        item = dataset[i]
+        assert isinstance(item, dict)
+        assert "question" in item
+        assert "answer" in item
+        assert isinstance(item["question"], str)
+        assert isinstance(item["answer"], str)
+
+
+def test_gsm_symbolic_iteration():
+    """Test that iteration respects dataset size"""
+    config = GSMSymbolicDatasetConfig(size=5, seed=42)  # Small size for testing
+    dataset = GSMSymbolicDataset(config)
+
+    # Test manual iteration
+    items = []
+    for item in dataset:
+        items.append(item)
+    assert len(items) == config.size, "Iterator should yield exactly size items"
+
+    # Test list conversion
+    items = list(dataset)
+    assert len(items) == config.size, "Iterator should yield exactly size items"
+
+    # Test multiple iterations
+    first_items = list(dataset)
+    second_items = list(dataset)
+    assert first_items == second_items, "Multiple iterations should yield same items"
+
+
+def test_gsm_symbolic_generators():
+    """Test generator loading and access"""
+    config = GSMSymbolicDatasetConfig()
+    dataset = GSMSymbolicDataset(config)
+
+    # Test lazy loading
+    assert dataset._generators is None
+    _ = dataset.generators  # Access to trigger loading
+    assert dataset._generators is not None
+
+    # Test generator mapping
+    assert isinstance(dataset.generators, dict)
+    assert len(dataset.generators) > 0
+    i = 0
+    rng = Random(18)
+    for key in sorted(dataset.generators.keys()):
+        generator = dataset.generators[key]
+        assert callable(generator)
+
+        print(i, key)
+        answer_set = set()
+        question_set = set()
+        for j in range(10):
+            x = generator(rng, difficulty=1.0)
+            question_set.add(x["question"])
+            answer_set.add(x["answer"])
+            # if j == 123:
+            #     print(f"[{j}] q: {x['question']}")
+            #     print(f"a: {x['answer']}")
+            #     print()
+        print(f"ok: q={len(question_set)}, a={len(answer_set)}")
+
+        i += 1