textpolicy 0.1.0__tar.gz → 0.1.2__tar.gz

{textpolicy-0.1.0/textpolicy.egg-info → textpolicy-0.1.2}/PKG-INFO

@@ -1,13 +1,23 @@
 Metadata-Version: 2.4
 Name: textpolicy
-Version: 0.1.0
-Summary: MLX-optimized reward and verification system for text generation RL
+Version: 0.1.2
+Summary: Reinforcement learning for text generation on MLX (Apple Silicon): GRPO/GSPO, environments, rollout, rewards, LoRA/QLoRA
+Project-URL: Homepage, https://github.com/teilomillet/textpolicy
+Project-URL: Repository, https://github.com/teilomillet/textpolicy
+Project-URL: Documentation, https://github.com/teilomillet/textpolicy#readme
+Project-URL: Changelog, https://github.com/teilomillet/textpolicy/blob/main/CHANGELOG.md
+Keywords: reinforcement-learning,text-generation,mlx,apple-silicon,lora,qlora,grpo,gspo,rlhf
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: MacOS
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=2.3.2
-Requires-Dist: mlx>=0.21.0
-Requires-Dist: mlx-lm>=0.21.0
+Requires-Dist: mlx>=0.22.0
+Requires-Dist: mlx-lm>=0.22.0
 Requires-Dist: gymnasium>=0.29.0
 Requires-Dist: psutil>=7.0.0
 Requires-Dist: wandb>=0.21.1

textpolicy-0.1.2/pyproject.toml

@@ -0,0 +1,56 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["textpolicy*"]
+
+[project]
+name = "textpolicy"
+version = "0.1.2"
+description = "Reinforcement learning for text generation on MLX (Apple Silicon): GRPO/GSPO, environments, rollout, rewards, LoRA/QLoRA"
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "numpy>=2.3.2",
+    "mlx>=0.22.0", # Core MLX framework for Apple Silicon acceleration (tested up to 0.30.5)
+    "mlx-lm>=0.22.0", # MLX language models for inference (tested up to 0.30.6)
+    "gymnasium>=0.29.0",
+    "psutil>=7.0.0",
+    "wandb>=0.21.1",
+    "aiohttp>=3.12.15",
+    "pytest>=8.4.1",
+]
+
+keywords = [
+    "reinforcement-learning",
+    "text-generation",
+    "mlx",
+    "apple-silicon",
+    "lora",
+    "qlora",
+    "grpo",
+    "gspo",
+    "rlhf",
+]
+
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: MacOS",
+    "Intended Audience :: Developers",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "License :: OSI Approved :: MIT License",
+]
+
+[project.urls]
+Homepage = "https://github.com/teilomillet/textpolicy"
+Repository = "https://github.com/teilomillet/textpolicy"
+Documentation = "https://github.com/teilomillet/textpolicy#readme"
+Changelog = "https://github.com/teilomillet/textpolicy/blob/main/CHANGELOG.md"
+
+[project.scripts]
+textpolicy = "textpolicy.cli:main"
+
+[project.optional-dependencies]
+external = ["aiohttp>=3.8.0", "pydantic>=2.0.0"]
+dev = ["pytest>=7.0.0", "black>=22.0.0", "ruff>=0.1.0"]

textpolicy-0.1.2/tests/test_issue_fixes.py

@@ -0,0 +1,218 @@
+"""
+Tests for GitHub Issues #2 and #3 fixes.
+
+Issue #2: TextGenerationEnv cannot pass example metadata to reward functions
+Issue #3: Export @verifier decorator at top level for API consistency
+"""
+
+import pytest
+
+
+class TestIssue3VerifierExport:
+    """Test that @verifier decorator is exported at top level (Issue #3)."""
+
+    def test_verifier_accessible_at_top_level(self):
+        """tp.verifier should be accessible without deep import."""
+        import textpolicy as tp
+
+        assert hasattr(tp, "verifier"), "verifier should be exported at top level"
+        assert callable(tp.verifier), "verifier should be callable"
+
+    def test_verifier_same_as_deep_import(self):
+        """tp.verifier should be the same function as deep import."""
+        import textpolicy as tp
+        from textpolicy.rewards import verifier
+
+        assert tp.verifier is verifier, "tp.verifier should be same as textpolicy.rewards.verifier"
+
+    def test_verifier_decorator_works(self):
+        """@tp.verifier should work as a decorator."""
+        import textpolicy as tp
+
+        @tp.verifier
+        def my_test_verifier(prompt, completion, example, **kwargs):
+            return len(completion) > 0
+
+        # Verify it's registered and callable
+        assert callable(my_test_verifier)
+        result = my_test_verifier("test prompt", "test completion", {})
+        assert isinstance(result, bool)
+
+    def test_reward_and_verifier_both_at_top_level(self):
+        """Both @reward and @verifier should be at top level for API consistency."""
+        import textpolicy as tp
+
+        assert hasattr(tp, "reward"), "reward should be at top level"
+        assert hasattr(tp, "verifier"), "verifier should be at top level"
+
+
+@pytest.mark.integration
+class TestIssue2ExamplesParameter:
+    """Test that TextGenerationEnv passes example metadata to reward functions (Issue #2)."""
+
+    @pytest.fixture
+    def dummy_tokenizer(self):
+        """Provide a minimal tokenizer for tests."""
+        class DummyTokenizer:
+            def encode(self, text):
+                return [ord(c) % 256 for c in text]
+
+            def decode(self, ids):
+                return "".join(chr(int(i) % 256) for i in ids)
+
+        return DummyTokenizer()
+
+    def test_env_accepts_examples_parameter(self, dummy_tokenizer):
+        """TextGenerationEnv should accept an examples parameter."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+
+        prompts = ["Hello", "World"]
+        examples = [{"key": "value1"}, {"key": "value2"}]
+
+        def reward_fn(prompt, completion, example, **kwargs):
+            return 1.0
+
+        # Should not raise
+        env = TextGenerationEnv(
+            prompts=prompts,
+            reward_fn=reward_fn,
+            tokenizer=dummy_tokenizer,
+            examples=examples,
+        )
+        assert env.examples == examples
+
+    def test_env_defaults_to_empty_dicts_when_no_examples(self, dummy_tokenizer):
+        """When examples not provided, should default to empty dicts."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+
+        prompts = ["Hello", "World"]
+
+        def reward_fn(prompt, completion, example, **kwargs):
+            return 1.0
+
+        env = TextGenerationEnv(
+            prompts=prompts, reward_fn=reward_fn, tokenizer=dummy_tokenizer
+        )
+
+        assert env.examples == [{}, {}]
+
+    def test_env_validates_examples_length(self, dummy_tokenizer):
+        """Should raise ValueError if examples length != prompts length."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+
+        prompts = ["Hello", "World"]
+        examples = [{"key": "value1"}]  # Wrong length
+
+        def reward_fn(prompt, completion, example, **kwargs):
+            return 1.0
+
+        with pytest.raises(ValueError, match="examples length.*must match prompts length"):
+            TextGenerationEnv(
+                prompts=prompts,
+                reward_fn=reward_fn,
+                tokenizer=dummy_tokenizer,
+                examples=examples,
+            )
+
+    def test_example_passed_to_reward_function(self, dummy_tokenizer):
+        """Reward function should receive the correct example for each prompt."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+
+        prompts = ["Question 1", "Question 2"]
+        examples = [
+            {"db_id": "database_1", "gold_sql": "SELECT 1"},
+            {"db_id": "database_2", "gold_sql": "SELECT 2"},
+        ]
+
+        received_examples = []
+
+        def capture_reward(prompt, completion, example, **kwargs):
+            received_examples.append(example.copy())
+            return 1.0
+
+        env = TextGenerationEnv(
+            prompts=prompts,
+            reward_fn=capture_reward,
+            tokenizer=dummy_tokenizer,
+            examples=examples,
+        )
+
+        # Episode 0 should use examples[0]
+        env.reset()
+        env.step("some response")
+        assert received_examples[0] == {"db_id": "database_1", "gold_sql": "SELECT 1"}
+
+        # Episode 1 should use examples[1]
+        env.reset()
+        env.step("another response")
+        assert received_examples[1] == {"db_id": "database_2", "gold_sql": "SELECT 2"}
+
+    def test_examples_cycle_with_prompts(self, dummy_tokenizer):
+        """Examples should cycle correctly when prompts cycle."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+
+        prompts = ["P1", "P2"]
+        examples = [{"idx": 0}, {"idx": 1}]
+
+        received_indices = []
+
+        def capture_reward(prompt, completion, example, **kwargs):
+            received_indices.append(example.get("idx"))
+            return 1.0
+
+        env = TextGenerationEnv(
+            prompts=prompts,
+            reward_fn=capture_reward,
+            tokenizer=dummy_tokenizer,
+            examples=examples,
+        )
+
+        # Run 4 episodes (should cycle through prompts twice)
+        for _ in range(4):
+            env.reset()
+            env.step("response")
+
+        # Should have received [0, 1, 0, 1]
+        assert received_indices == [0, 1, 0, 1]
+
+    def test_litmus_test_from_issue(self, dummy_tokenizer):
+        """Run the exact litmus test from Issue #2."""
+        from textpolicy.environment.text_generation import TextGenerationEnv
+        import textpolicy as tp
+
+        examples = [
+            {"db_id": "concert_singer", "gold_sql": "SELECT COUNT(*) FROM singer"},
+            {"db_id": "pets_1", "gold_sql": "SELECT COUNT(*) FROM pets"},
+        ]
+        prompts = [
+            "Schema: singer(id, name)\nQuestion: How many singers?",
+            "Schema: pets(id, name)\nQuestion: How many pets?",
+        ]
+
+        captured_db_ids = []
+
+        @tp.reward
+        def check_example(prompt, completion, example, **kwargs):
+            db_id = example.get("db_id")
+            captured_db_ids.append(db_id)
+            return 1.0
+
+        env = TextGenerationEnv(prompts, check_example, examples=examples, tokenizer=dummy_tokenizer)
+
+        # First episode
+        env.reset()
+        env.step("some action")
+
+        # Should have captured 'concert_singer'
+        assert captured_db_ids[0] == "concert_singer", f"Expected 'concert_singer', got {captured_db_ids[0]}"
+
+        # Second episode
+        env.reset()
+        env.step("another action")
+
+        # Should have captured 'pets_1'
+        assert captured_db_ids[1] == "pets_1", f"Expected 'pets_1', got {captured_db_ids[1]}"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

textpolicy-0.1.2/tests/test_mlx_compatibility.py

@@ -0,0 +1,315 @@
+"""
+MLX and mlx-lm Compatibility Tests
+
+These tests verify that textpolicy works correctly with MLX 0.30.x and mlx-lm 0.30.x.
+They cover the core APIs used throughout the codebase.
+
+Run with: pytest tests/test_mlx_compatibility.py -v
+"""
+
+import pytest
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@pytest.mark.unit
+class TestMLXCoreAPIs:
+    """Test core MLX APIs used in textpolicy."""
+
+    def test_mx_compile_decorator(self):
+        """Test @mx.compile decorator works (used in grpo.py, gspo.py, trainer.py)."""
+        @mx.compile
+        def compiled_fn(x, y):
+            return x + y
+
+        result = compiled_fn(mx.array([1.0]), mx.array([2.0]))
+        assert float(result[0]) == 3.0
+
+    def test_mx_compile_with_value_and_grad(self):
+        """Test mx.compile with nn.value_and_grad (used in trainer.py:90)."""
+        class SimpleModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(4, 2)
+            def __call__(self, x):
+                return self.linear(x)
+
+        model = SimpleModel()
+        mx.eval(model.parameters())
+
+        def loss_fn(model, x, y):
+            pred = model(x)
+            return mx.mean((pred - y) ** 2)
+
+        loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
+        x = mx.random.normal((2, 4))
+        y = mx.random.normal((2, 2))
+        loss, grads = loss_and_grad_fn(model, x, y)
+
+        assert not mx.isnan(loss)
+        assert isinstance(grads, dict)
+
+    def test_array_operations(self):
+        """Test array operations used in GRPO/GSPO algorithms."""
+        # Ratio computation (grpo.py, gspo.py)
+        old_lp = mx.array([-1.0, -1.2, -0.8])
+        new_lp = mx.array([-1.1, -1.0, -0.9])
+        ratios = mx.exp(new_lp - old_lp)
+
+        assert ratios.shape == (3,)
+        assert all(not mx.isnan(r) for r in ratios)
+
+        # Clipping (PPO-style)
+        clipped = mx.clip(ratios, 0.8, 1.2)
+        # Use small epsilon for float comparison
+        assert all(0.8 - 1e-6 <= float(c) <= 1.2 + 1e-6 for c in clipped)
+
+    def test_array_slicing_with_python_ints(self):
+        """Test array slicing with Python integers (used in gspo.py)."""
+        arr = mx.array([1.0, 2.0, 3.0, 4.0, 5.0])
+
+        # This pattern is used in compute_sequence_importance_weights
+        current_idx = 0
+        for seq_len in [2, 3]:
+            result = arr[current_idx:current_idx + seq_len]
+            assert result.shape[0] == seq_len
+            current_idx += seq_len
+
+
+@pytest.mark.unit
+class TestMLXLMAPIs:
+    """Test mlx-lm APIs used in textpolicy."""
+
+    def test_mlx_lm_imports(self):
+        """Test mlx_lm core imports (used in mlx_generation.py)."""
+        from mlx_lm import load, generate
+        assert callable(load)
+        assert callable(generate)
+
+    def test_sample_utils_imports(self):
+        """Test sample_utils imports (used in mlx_generation.py:27)."""
+        from mlx_lm.sample_utils import make_sampler, make_logits_processors
+        assert callable(make_sampler)
+        assert callable(make_logits_processors)
+
+    def test_make_sampler_signature(self):
+        """Test make_sampler accepts expected parameters."""
+        from mlx_lm.sample_utils import make_sampler
+
+        # These params are used in mlx_generation.py:77-81
+        sampler = make_sampler(
+            temp=0.7,
+            top_p=0.9,
+            min_p=0.0,
+            min_tokens_to_keep=2
+        )
+        assert sampler is not None
+
+    def test_quantize_model_signature(self):
+        """Test quantize_model has expected parameters (used in lora.py:335)."""
+        from mlx_lm.utils import quantize_model
+        import inspect
+
+        sig = inspect.signature(quantize_model)
+        params = list(sig.parameters.keys())
+
+        # These params are used in lora.py:342-348
+        expected = ['model', 'config', 'q_group_size', 'q_bits', 'quant_predicate']
+        for p in expected:
+            assert p in params, f"Expected param '{p}' not found in quantize_model"
+
+
+@pytest.mark.unit
+@pytest.mark.algorithm
+class TestGRPOWithMLX:
+    """Test GRPO algorithms with current MLX version."""
+
+    def test_compute_advantages(self):
+        """Test GRPO compute_advantages."""
+        from textpolicy.algorithms import grpo
+
+        rewards = mx.array([1.0, 0.5, -0.5, 0.8, 0.2])
+        advantages = grpo.compute_advantages(rewards)
+
+        assert advantages.shape == rewards.shape
+        # Group-relative: mean should be ~0
+        assert abs(float(mx.mean(advantages))) < 1e-5
+
+    def test_compute_advantages_compiled(self):
+        """Test compiled version of compute_advantages."""
+        from textpolicy.algorithms import grpo
+
+        rewards = mx.array([1.0, 0.5, -0.5, 0.8])
+        advantages = grpo.compute_advantages_compiled(rewards)
+
+        assert advantages.shape == rewards.shape
+
+    def test_policy_loss(self):
+        """Test GRPO policy_loss computation."""
+        from textpolicy.algorithms import grpo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        advantages = mx.array([0.6, 0.1, -0.9, 0.4, -0.2])
+
+        loss = grpo.policy_loss(old_lp, new_lp, advantages, clip_ratio=0.2)
+
+        assert not mx.isnan(loss)
+        assert loss.shape == ()  # Scalar
+
+    def test_entropy_bonus(self):
+        """Test entropy bonus computation."""
+        from textpolicy.algorithms import grpo
+
+        logprobs = mx.array([-1.0, -2.0, -0.5, -1.5])
+        entropy = grpo.entropy_bonus(logprobs, coefficient=0.01)
+
+        assert not mx.isnan(entropy)
+
+
+@pytest.mark.unit
+@pytest.mark.algorithm
+class TestGSPOWithMLX:
+    """Test GSPO algorithms with current MLX version."""
+
+    def test_compute_sequence_importance_weights(self):
+        """Test sequence-level importance weights."""
+        from textpolicy.algorithms import gspo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        seq_lens = [2, 3]
+
+        weights = gspo.compute_sequence_importance_weights(
+            old_lp, new_lp, seq_lens, clip_ratio=0.2
+        )
+
+        assert len(weights) == len(seq_lens)
+        assert all(not mx.isnan(w) for w in weights)
+
+    def test_compute_hybrid_importance_weights(self):
+        """Test hybrid importance weights."""
+        from textpolicy.algorithms import gspo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        seq_lens = [2, 3]
+
+        weights = gspo.compute_hybrid_importance_weights(
+            old_lp, new_lp, seq_lens, alpha=0.5, beta=0.5
+        )
+
+        assert len(weights) == sum(seq_lens)  # Token-level weights
+
+    def test_gspo_policy_loss_sequence_variant(self):
+        """Test GSPO policy loss with sequence variant."""
+        from textpolicy.algorithms import gspo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        seq_lens = [2, 3]
+        advantages = mx.array([0.5, -0.3])
+
+        loss = gspo.gspo_policy_loss(
+            old_lp, new_lp, advantages, seq_lens,
+            variant='sequence', clip_ratio=0.2
+        )
+
+        assert not mx.isnan(loss)
+        assert loss.shape == ()
+
+    def test_gspo_policy_loss_hybrid_variant(self):
+        """Test GSPO policy loss with hybrid variant."""
+        from textpolicy.algorithms import gspo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        seq_lens = [2, 3]
+        advantages = mx.array([0.5, -0.3])
+
+        loss = gspo.gspo_policy_loss(
+            old_lp, new_lp, advantages, seq_lens,
+            variant='hybrid', clip_ratio=0.2
+        )
+
+        assert not mx.isnan(loss)
+
+    def test_gspo_policy_loss_token_variant(self):
+        """Test GSPO policy loss with token variant (GRPO fallback)."""
+        from textpolicy.algorithms import gspo
+
+        old_lp = mx.array([-1.0, -1.2, -0.8, -1.1, -0.9])
+        new_lp = mx.array([-1.1, -1.0, -0.9, -1.0, -1.0])
+        seq_lens = [2, 3]
+        advantages = mx.array([0.5, -0.3])
+
+        loss = gspo.gspo_policy_loss(
+            old_lp, new_lp, advantages, seq_lens,
+            variant='token', clip_ratio=0.2
+        )
+
+        assert not mx.isnan(loss)
+
+
+@pytest.mark.unit
+class TestTrainerWithMLX:
+    """Test Trainer compilation with MLX."""
+
+    def test_trainer_compiles_loss_function(self):
+        """Test that Trainer correctly compiles the loss function."""
+        import mlx.optimizers as optim
+        from textpolicy.training import Trainer
+        from textpolicy.algorithms import grpo
+
+        class SimpleModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 5)
+            def __call__(self, x):
+                return self.linear(x)
+
+        model = SimpleModel()
+        mx.eval(model.parameters())
+
+        def loss_fn(model, batch):
+            obs = batch.get('observations', mx.zeros((1, 10)))
+            return mx.mean(model(obs) ** 2)
+
+        trainer = Trainer(
+            model=model,
+            loss_fn=loss_fn,
+            optimizer=optim.Adam(learning_rate=1e-3),
+            advantage_fn=grpo.compute_advantages,
+            compile_training=True
+        )
+
+        # Verify compilation happened
+        assert trainer.loss_and_grad_fn is not None
+
+
+@pytest.mark.unit
+class TestLoRAWithMLX:
+    """Test LoRA functions with MLX."""
+
+    def test_lora_functions_importable(self):
+        """Test all LoRA functions can be imported."""
+        from textpolicy.generation.lora import (
+            apply_lora,
+            freeze_base,
+            extract_params,
+            merge_weights,
+            create_lora_setup,
+            create_qlora_setup,
+            apply_quantization_to_model,
+            compute_lora_memory_savings
+        )
+
+        # All should be callable
+        assert callable(apply_lora)
+        assert callable(freeze_base)
+        assert callable(extract_params)
+        assert callable(merge_weights)
+        assert callable(create_lora_setup)
+        assert callable(create_qlora_setup)
+        assert callable(apply_quantization_to_model)
+        assert callable(compute_lora_memory_savings)

textpolicy-0.1.2/tests/test_training_pipeline.py

@@ -0,0 +1,240 @@
+"""
+Training Pipeline Integration Tests
+
+These tests verify that the complete training pipeline works correctly,
+including gradient computation and parameter updates.
+"""
+
+import pytest
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optim
+
+
+@pytest.mark.integration
+class TestTrainingPipeline:
+    """Test the complete training pipeline with gradient updates."""
+
+    def test_grpo_training_step_produces_finite_loss(self):
+        """Test that GRPO training produces finite loss values."""
+        from textpolicy.algorithms import grpo
+
+        # Create simple model
+        class TinyLM(nn.Module):
+            def __init__(self, vocab_size=100, hidden=32):
+                super().__init__()
+                self.embed = nn.Embedding(vocab_size, hidden)
+                self.proj = nn.Linear(hidden, vocab_size)
+
+            def __call__(self, x):
+                return self.proj(self.embed(x))
+
+        model = TinyLM()
+        mx.eval(model.parameters())
+
+        # Prepare batch data (mimics what data_selector returns)
+        old_logprobs = mx.array([-2.5, -3.1, -2.8, -2.9, -3.0])
+        observations = mx.array([10, 20, 30, 40, 50])
+        actions = mx.array([15, 25, 35, 45, 55])
+        rewards = mx.array([1.0, 0.5, -0.5, 0.8, 0.2])
+
+        # Compute new logprobs from model
+        logits = model(observations)
+        log_probs = nn.log_softmax(logits, axis=-1)
+        new_logprobs = mx.take_along_axis(log_probs, actions[:, None], axis=-1).squeeze(-1)
+
+        # Compute GRPO loss
+        advantages = grpo.compute_advantages(rewards)
+        loss = grpo.policy_loss(old_logprobs, new_logprobs, advantages, clip_ratio=0.2)
+
+        mx.eval(loss)
+        assert not mx.isnan(loss), "Loss should not be NaN"
+        assert not mx.isinf(loss), "Loss should not be infinite"
+
+    def test_grpo_gradients_flow(self):
+        """Test that gradients flow through the GRPO loss."""
+        from textpolicy.algorithms import grpo
+
+        class TinyLM(nn.Module):
+            def __init__(self, vocab_size=100, hidden=32):
+                super().__init__()
+                self.embed = nn.Embedding(vocab_size, hidden)
+                self.proj = nn.Linear(hidden, vocab_size)
+
+            def __call__(self, x):
+                return self.proj(self.embed(x))
+
+        model = TinyLM()
+        mx.eval(model.parameters())
+
+        def loss_fn(model):
+            observations = mx.array([10, 20, 30, 40, 50])
+            actions = mx.array([15, 25, 35, 45, 55])
+            old_logprobs = mx.array([-2.5, -3.1, -2.8, -2.9, -3.0])
+            rewards = mx.array([1.0, 0.5, -0.5, 0.8, 0.2])
+
+            logits = model(observations)
+            log_probs = nn.log_softmax(logits, axis=-1)
+            new_logprobs = mx.take_along_axis(log_probs, actions[:, None], axis=-1).squeeze(-1)
+
+            advantages = grpo.compute_advantages(rewards)
+            return grpo.policy_loss(old_logprobs, new_logprobs, advantages, clip_ratio=0.2)
+
+        # Compute gradients
+        loss_and_grad = nn.value_and_grad(model, loss_fn)
+        loss, grads = loss_and_grad(model)
+
+        mx.eval(loss, grads)
+
+        assert not mx.isnan(loss), "Loss should not be NaN"
+
+        # Check that at least some gradients are non-zero
+        has_nonzero_grad = False
+        for name, grad in grads.items():
+            if isinstance(grad, dict):
+                for subname, subgrad in grad.items():
+                    if mx.any(subgrad != 0):
+                        has_nonzero_grad = True
+                        break
+            elif mx.any(grad != 0):
+                has_nonzero_grad = True
+                break
+
+        assert has_nonzero_grad, "Should have at least some non-zero gradients"
+
+    def test_optimizer_updates_parameters(self):
+        """Test that optimizer actually updates model parameters."""
+        from textpolicy.algorithms import grpo
+
+        class TinyLM(nn.Module):
+            def __init__(self, vocab_size=100, hidden=32):
+                super().__init__()
+                self.embed = nn.Embedding(vocab_size, hidden)
+                self.proj = nn.Linear(hidden, vocab_size)
+
+            def __call__(self, x):
+                return self.proj(self.embed(x))
+
+        model = TinyLM()
+        optimizer = optim.Adam(learning_rate=0.01)
+        mx.eval(model.parameters())
+
+        # Store initial parameter values
+        initial_weight = mx.array(model.proj.weight)  # Create a copy
+
+        def loss_fn(model):
+            observations = mx.array([10, 20, 30])
+            actions = mx.array([15, 25, 35])
+            old_logprobs = mx.array([-2.5, -3.1, -2.8])
+            rewards = mx.array([1.0, 0.5, -0.5])
+
+            logits = model(observations)
+            log_probs = nn.log_softmax(logits, axis=-1)
+            new_logprobs = mx.take_along_axis(log_probs, actions[:, None], axis=-1).squeeze(-1)
+
+            advantages = grpo.compute_advantages(rewards)
+            return grpo.policy_loss(old_logprobs, new_logprobs, advantages)
+
+        loss_and_grad = nn.value_and_grad(model, loss_fn)
+
+        # Run a few training steps
+        for _ in range(3):
+            loss, grads = loss_and_grad(model)
+            optimizer.update(model, grads)
+            mx.eval(model.parameters())
+
+        # Check parameters changed
+        final_weight = model.proj.weight
+        params_changed = not mx.allclose(initial_weight, final_weight)
+
+        assert params_changed, "Parameters should change after optimization"
+
+    def test_gspo_training_step_produces_finite_loss(self):
+        """Test that GSPO training produces finite loss values."""
+        from textpolicy.algorithms import gspo
+
+        class TinyLM(nn.Module):
+            def __init__(self, vocab_size=100, hidden=32):
+                super().__init__()
+                self.embed = nn.Embedding(vocab_size, hidden)
+                self.proj = nn.Linear(hidden, vocab_size)
+
+            def __call__(self, x):
+                return self.proj(self.embed(x))
+
+        model = TinyLM()
+        mx.eval(model.parameters())
+
+        # Prepare batch data
+        old_logprobs = mx.array([-2.5, -3.1, -2.8, -2.9, -3.0])
+        observations = mx.array([10, 20, 30, 40, 50])
+        actions = mx.array([15, 25, 35, 45, 55])
+
+        # Compute new logprobs from model
+        logits = model(observations)
+        log_probs = nn.log_softmax(logits, axis=-1)
+        new_logprobs = mx.take_along_axis(log_probs, actions[:, None], axis=-1).squeeze(-1)
+
+        # GSPO needs sequence-level advantages
+        sequence_lengths = [2, 3]  # 2 sequences
+        advantages = mx.array([0.5, -0.3])  # Per-sequence
+
+        # Test all GSPO variants
+        for variant in ['sequence', 'hybrid', 'token']:
+            loss = gspo.gspo_policy_loss(
+                old_logprobs, new_logprobs, advantages, sequence_lengths,
+                variant=variant, clip_ratio=0.2
+            )
+            mx.eval(loss)
+            assert not mx.isnan(loss), f"GSPO {variant} loss should not be NaN"
+            assert not mx.isinf(loss), f"GSPO {variant} loss should not be infinite"
+
+
+@pytest.mark.integration
+class TestCompiledTraining:
+    """Test that @mx.compile works correctly with training."""
+
+    def test_compiled_function_works(self):
+        """Test that @mx.compile decorator works with loss functions."""
+        @mx.compile
+        def compiled_loss(x, y):
+            return mx.mean((x - y) ** 2)
+
+        x = mx.array([1.0, 2.0, 3.0])
+        y = mx.array([1.1, 2.1, 3.1])
+
+        loss = compiled_loss(x, y)
+        mx.eval(loss)
+
+        assert not mx.isnan(loss), "Compiled loss should not be NaN"
+        assert float(loss) > 0, "Loss should be positive"
+
+    def test_trainer_compiled_loss_function(self):
+        """Test that Trainer's compiled loss function works correctly.
+
+        This mimics how Trainer.py:90 compiles the loss function.
+        """
+        from textpolicy.training import Trainer
+        from textpolicy.algorithms import grpo
+
+        class TinyLM(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.linear = nn.Linear(10, 10)
+            def __call__(self, x):
+                return self.linear(x)
+
+        model = TinyLM()
+        mx.eval(model.parameters())
+
+        # Create trainer with compilation enabled
+        trainer = Trainer(
+            model=model,
+            loss_fn=grpo.policy_loss,
+            optimizer=optim.Adam(learning_rate=1e-3),
+            advantage_fn=grpo.compute_advantages,
+            compile_training=True
+        )
+
+        # Verify the compiled function was created
+        assert trainer.loss_and_grad_fn is not None, "Compiled loss function should exist"

{textpolicy-0.1.0 → textpolicy-0.1.2}/textpolicy/__init__.py

@@ -30,7 +30,7 @@ from .validate import validate_installation
 
 # Export core reward functions and the reward decorator
 from .rewards.basic import length_reward, keyword_reward, perplexity_reward, accuracy_reward
-from .rewards.registry import reward
+from .rewards.registry import reward, verifier
 
 # Build __all__ combining submodule __all__ lists and additional symbols
 __all__ = (
@@ -48,5 +48,6 @@ __all__ = (
         "perplexity_reward",
         "accuracy_reward",
         "reward",
+        "verifier",
     ]
 )

{textpolicy-0.1.0 → textpolicy-0.1.2}/textpolicy/environment/text_generation.py

@@ -630,24 +630,31 @@ class TextGenerationEnv(Environment):
         reward_fn: Callable[[str, str, dict], float],
         max_tokens: int = 25,
         seed: int = 42,
-        tokenizer: Any = None
+        tokenizer: Any = None,
+        examples: Optional[List[dict]] = None
     ):
         """
         Initialize simple text generation environment.
-        
+
         Args:
             prompts: List of prompts to cycle through
             reward_fn: Function that computes reward from (prompt, completion, example)
             max_tokens: Maximum tokens to generate per response
             seed: Random seed for reproducible behavior
             tokenizer: Tokenizer for converting prompts to tokens (required for MLX compatibility)
+            examples: Optional list of example dicts to pass to reward function. If provided,
+                      must have same length as prompts. examples[i] is passed when prompts[i] is used.
         """
         super().__init__()
         
         if tokenizer is None:
             raise ValueError("tokenizer is required for TextGenerationEnv to work with MLX rollout system")
-        
+
+        if examples is not None and len(examples) != len(prompts):
+            raise ValueError(f"examples length ({len(examples)}) must match prompts length ({len(prompts)})")
+
         self.prompts = prompts
+        self.examples = examples if examples is not None else [{} for _ in prompts]
         self.reward_fn = reward_fn
         self.max_tokens = max_tokens
         self.tokenizer = tokenizer
@@ -735,10 +742,11 @@ class TextGenerationEnv(Environment):
         
         # Compute reward using provided reward function
         # Pass tokenizer for EOS token detection and truncation detection
+        prompt_index = self.current_episode % len(self.prompts)
         reward = self.reward_fn(
             prompt=self.current_prompt,
             completion=response_text,
-            example={},
+            example=self.examples[prompt_index],
             tokenizer=self.tokenizer,  # Pass tokenizer for EOS detection
             truncated=truncated        # Pass truncation flag from environment
         )

{textpolicy-0.1.0 → textpolicy-0.1.2/textpolicy.egg-info}/PKG-INFO

@@ -1,13 +1,23 @@
 Metadata-Version: 2.4
 Name: textpolicy
-Version: 0.1.0
-Summary: MLX-optimized reward and verification system for text generation RL
+Version: 0.1.2
+Summary: Reinforcement learning for text generation on MLX (Apple Silicon): GRPO/GSPO, environments, rollout, rewards, LoRA/QLoRA
+Project-URL: Homepage, https://github.com/teilomillet/textpolicy
+Project-URL: Repository, https://github.com/teilomillet/textpolicy
+Project-URL: Documentation, https://github.com/teilomillet/textpolicy#readme
+Project-URL: Changelog, https://github.com/teilomillet/textpolicy/blob/main/CHANGELOG.md
+Keywords: reinforcement-learning,text-generation,mlx,apple-silicon,lora,qlora,grpo,gspo,rlhf
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: MacOS
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
 Requires-Python: >=3.12
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy>=2.3.2
-Requires-Dist: mlx>=0.21.0
-Requires-Dist: mlx-lm>=0.21.0
+Requires-Dist: mlx>=0.22.0
+Requires-Dist: mlx-lm>=0.22.0
 Requires-Dist: gymnasium>=0.29.0
 Requires-Dist: psutil>=7.0.0
 Requires-Dist: wandb>=0.21.1

{textpolicy-0.1.0 → textpolicy-0.1.2}/textpolicy.egg-info/SOURCES.txt

@@ -3,9 +3,12 @@ README.md
 pyproject.toml
 tests/test_gspo_verification.py
 tests/test_integration_e2e_training.py
+tests/test_issue_fixes.py
+tests/test_mlx_compatibility.py
 tests/test_reward_signatures.py
 tests/test_rollout_rewards.py
 tests/test_runner_step_enforcement.py
+tests/test_training_pipeline.py
 tests/test_validate_installation.py
 textpolicy/__init__.py
 textpolicy/__main__.py

{textpolicy-0.1.0 → textpolicy-0.1.2}/textpolicy.egg-info/requires.txt

@@ -1,6 +1,6 @@
 numpy>=2.3.2
-mlx>=0.21.0
-mlx-lm>=0.21.0
+mlx>=0.22.0
+mlx-lm>=0.22.0
 gymnasium>=0.29.0
 psutil>=7.0.0
 wandb>=0.21.1

textpolicy-0.1.0/pyproject.toml

@@ -1,30 +0,0 @@
-[build-system]
-requires = ["setuptools>=61.0", "wheel"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools.packages.find]
-include = ["textpolicy*"]
-
-[project]
-name = "textpolicy"
-version = "0.1.0"
-description = "MLX-optimized reward and verification system for text generation RL"
-readme = "README.md"
-requires-python = ">=3.12"
-dependencies = [
-    "numpy>=2.3.2",
-    "mlx>=0.21.0", # Core MLX framework for Apple Silicon acceleration
-    "mlx-lm>=0.21.0", # MLX language models for inference
-    "gymnasium>=0.29.0", 
-    "psutil>=7.0.0",
-    "wandb>=0.21.1",
-    "aiohttp>=3.12.15",
-    "pytest>=8.4.1",
-]
-
-[project.scripts]
-textpolicy = "textpolicy.cli:main"
-
-[project.optional-dependencies]
-external = ["aiohttp>=3.8.0", "pydantic>=2.0.0"]
-dev = ["pytest>=7.0.0", "black>=22.0.0", "ruff>=0.1.0"]