textpolicy 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

textpolicy/generation/reload.py

@@ -0,0 +1,253 @@
+# textpolicy/generation/reload.py
+"""
+Pure functions for handling LoRA model reloading after training updates.
+
+Critical for RL training: After each training step, the LoRA adapters are updated.
+For the next rollout generation, we need to ensure the policy uses the updated model.
+
+Following our design principles:
+- Pure functions only
+- No state management 
+- Zero abstraction cost
+"""
+
+from typing import Dict, Tuple, Any, Optional, Callable
+import mlx.core as mx # type: ignore
+import mlx.nn as nn # type: ignore
+from .mlx_generation import load_model# type: ignore
+
+
+def save_adapters(
+    model: nn.Module,
+    adapter_path: str
+) -> None:
+    """
+    Pure function to save LoRA adapter weights.
+    
+    After training updates, save only the LoRA parameters to disk.
+    This is much faster than saving the full model.
+    
+    Args:
+        model: Model with trained LoRA adapters
+        adapter_path: Path to save adapters
+    """
+    # Extract LoRA parameters
+    lora_params = {}
+    for name, param in model.named_parameters():
+        if 'lora_' in name.lower() and param.requires_grad:
+            lora_params[name] = param
+    
+    # Save using MLX
+    mx.save_safetensors(adapter_path, lora_params)
+    print(f"✓ Saved LoRA adapters to {adapter_path}")
+
+
+def reload_model(
+    base_model_path: str,
+    adapter_path: str,
+    tokenizer_config: Optional[Dict] = None
+) -> Tuple[nn.Module, Any]:
+    """
+    Pure function to reload model with updated LoRA adapters.
+    
+    This is called after training to get an updated model for the next rollout.
+    Much more efficient than reloading the full model.
+    
+    Args:
+        base_model_path: Path to base model
+        adapter_path: Path to updated LoRA adapters
+        tokenizer_config: Optional tokenizer config
+        
+    Returns:
+        (updated_model, tokenizer): Model with updated adapters
+    """
+    # Load fresh model with updated adapters
+    model, tokenizer = load_model(
+        model_path=base_model_path,
+        adapter_path=adapter_path,
+        tokenizer_config=tokenizer_config
+    )
+    
+    print("✓ Reloaded model with updated adapters")
+    return model, tokenizer
+
+
+def create_reloadable_policy(
+    base_model_path: str,
+    initial_adapter_path: Optional[str],
+    generation_params: Dict[str, Any]
+) -> Tuple[Callable, Callable]:
+    """
+    Create a policy function that can be reloaded with updated LoRA adapters.
+    
+    Returns both the policy function and a reload function for efficiency.
+    
+    Args:
+        base_model_path: Path to base model
+        initial_adapter_path: Initial LoRA adapter path (can be None)
+        generation_params: Generation parameters
+        
+    Returns:
+        (policy_fn, reload_fn): Policy function and reload function
+    """
+    # Load initial model
+    current_model, tokenizer = load_model(
+        model_path=base_model_path,
+        adapter_path=initial_adapter_path
+    )
+    
+    # Store current state
+    current_state = {
+        'model': current_model,
+        'tokenizer': tokenizer,
+        'base_path': base_model_path,
+        'generation_params': generation_params
+    }
+    
+    def policy_fn(obs: mx.array, deterministic: bool = False) -> Tuple[mx.array, Dict[str, mx.array]]:
+        """Policy function that always uses the current model."""
+        from .mlx_generation import generate_tokens
+        
+        # Use current model (updated after each reload)
+        model = current_state['model']
+        tokenizer = current_state['tokenizer']
+        
+        # Adjust temperature
+        temp = 0.0 if deterministic else generation_params.get("temperature", 0.8)
+        
+        # Generate tokens
+        response_tokens, response_info = generate_tokens(
+            model=model,
+            tokenizer=tokenizer,
+            prompt_tokens=obs,
+            max_tokens=generation_params.get("max_tokens", 50),
+            temperature=temp,
+            top_p=generation_params.get("top_p", 0.95)
+        )
+        
+        # Extract logprobs from response info
+        logprobs = response_info.get('logprob', mx.array([]))
+        
+        extras = {
+            "logprob": logprobs,
+            "entropy": mx.mean(logprobs) if len(logprobs) > 0 else mx.array(0.0)
+        }
+        
+        return response_tokens, extras
+    
+    def reload_fn(adapter_path: str) -> None:
+        """Reload the model with updated adapters."""
+        updated_model, updated_tokenizer = reload_model(
+            base_model_path=current_state['base_path'],
+            adapter_path=adapter_path
+        )
+        
+        # Update current state
+        current_state['model'] = updated_model
+        current_state['tokenizer'] = updated_tokenizer
+        
+        print(f"✓ Policy reloaded with adapters from {adapter_path}")
+    
+    return policy_fn, reload_fn
+
+
+def create_training_loop_with_reload(
+    base_model_path: str,
+    adapter_save_path: str,
+    generation_params: Dict[str, Any],
+    trainer
+) -> Tuple[Callable, Callable]:
+    """
+    Create a complete training loop that handles LoRA reloading.
+    
+    This solves the LoRA update problem by automatically saving and reloading
+    adapters after each training step.
+    
+    Args:
+        base_model_path: Path to base model
+        adapter_save_path: Path to save LoRA adapters
+        generation_params: Generation parameters
+        trainer: GRPO trainer instance
+        
+    Returns:
+        (train_and_reload_fn, policy_fn): Training function and policy
+    """
+    # Create reloadable policy
+    policy_fn, reload_fn = create_reloadable_policy(
+        base_model_path=base_model_path,
+        initial_adapter_path=None,  # Start with base model
+        generation_params=generation_params
+    )
+    
+    def train_and_reload_fn(rollout_data) -> Dict[str, float]:
+        """
+        Train the model and reload policy with updated adapters.
+        
+        Args:
+            rollout_data: Rollout data for training
+            
+        Returns:
+            Training metrics
+        """
+        # 1. Train the model (updates LoRA adapters in-place)
+        metrics = trainer.train(rollout_data)
+        
+        # 2. Save updated LoRA adapters
+        save_adapters(trainer.model, adapter_save_path)
+        
+        # 3. Reload policy with updated adapters
+        reload_fn(adapter_save_path)
+        
+        print("✓ Training step complete, policy updated")
+        return metrics
+    
+    return train_and_reload_fn, policy_fn
+
+
+# Simple wrapper for the most common use case
+def create_auto_reload_setup(
+    model_path: str,
+    adapter_save_path: str = "./lora_adapters.safetensors",
+    **generation_params
+) -> Tuple[Callable, Callable, nn.Module, Any]:
+    """
+    Create complete auto-reloading setup for LoRA training.
+    
+    This is the main function most users should call.
+    
+    Args:
+        model_path: Path to base model
+        adapter_save_path: Where to save LoRA adapters
+        **generation_params: Generation parameters
+        
+    Returns:
+        (policy_fn, reload_fn, model, tokenizer): Complete setup
+    """
+    from .lora import create_lora_setup
+    
+    # Load and setup LoRA model
+    base_model, tokenizer = load_model(model_path)
+    
+    lora_config = {
+        "lora_layers": generation_params.pop("lora_layers", 8),
+        "lora_rank": generation_params.pop("lora_rank", 8),
+        "lora_scale": generation_params.pop("lora_scale", 20.0)
+    }
+    
+    lora_model, memory_stats = create_lora_setup(
+        model=base_model,
+        lora_config=lora_config
+    )
+    
+    # Create reloadable policy
+    policy_fn, reload_fn = create_reloadable_policy(
+        base_model_path=model_path,
+        initial_adapter_path=None,
+        generation_params=generation_params
+    )
+    
+    print("✓ Auto-reload setup complete")
+    print(f"   Memory savings: {memory_stats['memory_savings_percent']:.1f}%")
+    print(f"   Adapter save path: {adapter_save_path}")
+    
+    return policy_fn, reload_fn, lora_model, tokenizer

textpolicy/rewards/__init__.py

@@ -0,0 +1,137 @@
+# textpolicy/rewards/__init__.py
+"""
+Unified reward and verification system for MLX-optimized text generation.
+
+Key principles:
+- Decorator-based registration (@reward, @verifier)
+- Pure function composition with zero abstraction cost
+- MLX compilation for Apple Silicon optimization
+- Modular system allowing custom rewards and verifiers
+- Pre-filtering verification approach
+- Signature consistency: (prompt, completion, example, **kwargs)
+
+This system provides modular reward computation with MLX optimization.
+"""
+
+# Registry system
+from .registry import (
+    reward, verifier,  # Decorators for registration
+    RewardFunction, VerifierFunction, RewardConfig,
+    get_reward_function, get_verifier_function,
+    apply_verifiers_to_reward,
+    create_configured_reward_function,
+    list_registered_functions, clear_registries,
+    REWARD_REGISTRY, VERIFIER_REGISTRY
+)
+
+# Core reward functions (auto-registered via decorators)
+from .basic import (
+    length_reward,
+    keyword_reward, 
+    perplexity_reward,
+    accuracy_reward
+)
+
+# Core verifier functions (auto-registered via decorators)
+from .verifiers import (
+    length_verifier,
+    toxicity_verifier,
+    coherence_verifier,
+    factual_verifier,
+    has_greeting,
+    no_empty_response,
+    contains_keywords,
+    # Legacy compatibility (deprecated)
+    create_default_verifier_pipeline,
+    create_custom_verifier_pipeline
+)
+
+# MLX-optimized batch processing (following DESIGN_GUIDELINES.md)
+from .mlx_batch_processor import (
+    create_batch_reward_processor,
+    create_mlx_optimized_batch_processor,
+    create_async_batch_processor,
+    create_processor_from_config,
+    list_available_processors,
+    compute_length_rewards_vectorized,
+    compute_keyword_rewards_vectorized
+)
+
+# Core adapters and utilities for MLX reward system
+from .adapters import (
+    # Configuration models
+    MLXRewardConfig, MLXRewardSystemConfig,
+    # Sample types
+    MLXSample, MLXExternalRewardModel,
+    # Math utilities
+    extract_boxed_answer, normalize_math_answer, compute_f1_score,
+    # Bridge functions
+    create_mlx_batch_adapter, create_mlx_system_from_config,
+    samples_to_mlx_format, mlx_format_to_samples,
+    get_available_adapters,
+    # Reward functions
+    math_accuracy_reward, f1_score_reward
+)
+
+# Legacy systems (maintained for backward compatibility)
+from .rollout_rewards import (
+    RolloutRewardProcessor,
+    create_rollout_reward_processor,
+    process_episode_batch_rewards,
+    compute_reward_vector
+)
+
+from .integrated_system import (
+    IntegratedRewardConfig,
+    IntegratedRolloutRewardSystem,
+    create_integrated_reward_system,
+    process_episodes_with_quality_control,
+    compute_integrated_rewards
+)
+
+__all__ = [
+    # Registry system (primary interface)
+    "reward", "verifier",  # Decorators
+    "RewardFunction", "VerifierFunction", "RewardConfig",
+    "get_reward_function", "get_verifier_function",
+    "apply_verifiers_to_reward", "create_configured_reward_function",
+    "list_registered_functions", "clear_registries",
+    "REWARD_REGISTRY", "VERIFIER_REGISTRY",
+    
+    # Core reward functions (auto-registered)
+    "length_reward", "keyword_reward", "perplexity_reward", "accuracy_reward",
+    
+    # Core verifier functions (auto-registered)
+    "length_verifier", "toxicity_verifier", "coherence_verifier", "factual_verifier",
+    "has_greeting", "no_empty_response", "contains_keywords",
+    
+    # MLX-optimized batch processing (primary interface for training)
+    "create_batch_reward_processor",
+    "create_mlx_optimized_batch_processor", 
+    "create_async_batch_processor",
+    "create_processor_from_config",
+    "list_available_processors",
+    "compute_length_rewards_vectorized",
+    "compute_keyword_rewards_vectorized",
+    
+    # Core adapters and utilities
+    "MLXRewardConfig", "MLXRewardSystemConfig",
+    "MLXSample", "MLXExternalRewardModel", 
+    "extract_boxed_answer", "normalize_math_answer", "compute_f1_score",
+    "create_mlx_batch_adapter", "create_mlx_system_from_config",
+    "samples_to_mlx_format", "mlx_format_to_samples", "get_available_adapters",
+    "math_accuracy_reward", "f1_score_reward",
+    
+    # Legacy systems (backward compatibility)
+    "RolloutRewardProcessor", "create_rollout_reward_processor",
+    "process_episode_batch_rewards", "compute_reward_vector",
+    "IntegratedRewardConfig", "IntegratedRolloutRewardSystem",
+    "create_integrated_reward_system", "process_episodes_with_quality_control",
+    "compute_integrated_rewards",
+    "create_default_verifier_pipeline", "create_custom_verifier_pipeline",
+]
+
+# Import basic rewards and verifiers to trigger registration
+import textpolicy.rewards.basic
+import textpolicy.rewards.verifiers
+import textpolicy.rewards.adapters  # Register adapted rewards