textpolicy 0.1.1__tar.gz → 0.1.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: textpolicy
-Version: 0.1.1
+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
@@ -16,8 +16,8 @@ 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.1 → textpolicy-0.1.2}/pyproject.toml

@@ -7,15 +7,15 @@ include = ["textpolicy*"]
 
 [project]
 name = "textpolicy"
-version = "0.1.1"
+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.21.0", # Core MLX framework for Apple Silicon acceleration
-    "mlx-lm>=0.21.0", # MLX language models for inference
-    "gymnasium>=0.29.0", 
+    "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",

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.1 → textpolicy-0.1.2/textpolicy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: textpolicy
-Version: 0.1.1
+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
@@ -16,8 +16,8 @@ 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.1 → textpolicy-0.1.2}/textpolicy.egg-info/SOURCES.txt

@@ -4,9 +4,11 @@ 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.1 → 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