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

textpolicy/rewards/registry.py

@@ -0,0 +1,293 @@
+# textpolicy/rewards/registry.py
+"""
+Unified reward and verifier registry system following retrain's philosophy.
+
+This module provides decorator-based registration for rewards and verifiers,
+maintaining compatibility with MLX optimization and pure function composition.
+
+Key principles:
+- Decorator-based registration (@reward, @verifier)
+- Function signature consistency: (prompt, completion, example, **kwargs)
+- Pre-filtering verification approach
+- Global registries for modularity
+- MLX compilation support
+"""
+
+from typing import Callable, Dict, Any, List, Optional, Union
+import inspect
+import functools
+import mlx.core as mx
+from dataclasses import dataclass
+
+# Type definitions following retrain's patterns
+RewardFunction = Callable[[str, str, Dict[str, Any]], float]
+VerifierFunction = Callable[[str, str, Dict[str, Any]], bool]
+
+# Global registries following retrain's architecture
+REWARD_REGISTRY: Dict[str, RewardFunction] = {}
+VERIFIER_REGISTRY: Dict[str, VerifierFunction] = {}
+
+# Simple logging
+import logging
+logger = logging.getLogger(__name__)
+
+
+def reward(_func: Optional[RewardFunction] = None, *, name: Optional[str] = None) -> Union[Callable[[RewardFunction], RewardFunction], RewardFunction]:
+    """
+    Decorator to register reward functions following retrain's pattern.
+    
+    Usage:
+        @reward
+        def my_reward(prompt: str, completion: str, example: Dict[str, Any]) -> float:
+            return 1.0
+            
+        @reward(name="custom_name")
+        def another_reward(prompt: str, completion: str, example: Dict[str, Any]) -> float:
+            return 0.5
+    """
+    def decorator_reward(func: RewardFunction) -> RewardFunction:
+        if not callable(func):
+            raise TypeError(f"Object {getattr(func, '__name__', '<unknown>')} must be callable to be registered as reward.")
+        
+        registration_name = name if name is not None else func.__name__
+        
+        # Validate function signature for consistency
+        sig = inspect.signature(func)
+        expected_params = ['prompt', 'completion', 'example']
+        
+        if len(sig.parameters) < 3:
+            logger.warning(f"Reward function '{registration_name}' has fewer than 3 expected parameters. Ensure signature compatibility.")
+        
+        param_names = list(sig.parameters.keys())
+        for i, expected in enumerate(expected_params):
+            if i < len(param_names) and param_names[i] != expected:
+                logger.warning(f"Reward function '{registration_name}' parameter {i} is '{param_names[i]}', expected '{expected}'.")
+        
+        if registration_name in REWARD_REGISTRY:
+            logger.warning(f"Reward function '{registration_name}' already registered. Overwriting.")
+        
+        REWARD_REGISTRY[registration_name] = func
+        logger.info(f"Registered reward function: '{registration_name}' -> {func.__name__}")
+        
+        return func
+    
+    if _func is None:
+        # Called with parentheses: @reward() or @reward(name=...)
+        return decorator_reward
+    elif callable(_func):
+        # Called without parentheses: @reward
+        if name is not None:
+            raise TypeError("Cannot specify 'name' when using @reward without parentheses. Use @reward(name='...') instead.")
+        return decorator_reward(_func)
+    else:
+        raise TypeError("Invalid arguments supplied to @reward decorator.")
+
+
+def verifier(_func: Optional[VerifierFunction] = None, *, name: Optional[str] = None) -> Union[Callable[[VerifierFunction], VerifierFunction], VerifierFunction]:
+    """
+    Decorator to register verifier functions following retrain's pattern.
+    
+    Usage:
+        @verifier
+        def has_greeting(prompt: str, completion: str, example: Dict[str, Any]) -> bool:
+            return completion.lower().startswith("hello")
+            
+        @verifier(name="custom_check")
+        def custom_verifier(prompt: str, completion: str, example: Dict[str, Any]) -> bool:
+            return len(completion) > 10
+    """
+    def decorator_verifier(func: VerifierFunction) -> VerifierFunction:
+        if not callable(func):
+            raise TypeError(f"Object {getattr(func, '__name__', '<unknown>')} must be callable to be registered as verifier.")
+        
+        registration_name = name if name is not None else func.__name__
+        
+        # Validate function signature
+        sig = inspect.signature(func)
+        if len(sig.parameters) < 3:
+            logger.warning(f"Verifier function '{registration_name}' has fewer than 3 expected parameters (prompt, completion, example).")
+        
+        if registration_name in VERIFIER_REGISTRY:
+            logger.warning(f"Verifier function '{registration_name}' already registered. Overwriting.")
+        
+        VERIFIER_REGISTRY[registration_name] = func
+        logger.info(f"Registered verifier function: '{registration_name}' -> {func.__name__}")
+        
+        return func
+    
+    if _func is None:
+        return decorator_verifier
+    elif callable(_func):
+        if name is not None:
+            raise TypeError("Cannot specify 'name' when using @verifier without parentheses. Use @verifier(name='...') instead.")
+        return decorator_verifier(_func)
+    else:
+        raise TypeError("Invalid arguments supplied to @verifier decorator.")
+
+
+def get_reward_function(name: str) -> Optional[RewardFunction]:
+    """Retrieve a registered reward function by name."""
+    func = REWARD_REGISTRY.get(name)
+    if func is None:
+        available = list(REWARD_REGISTRY.keys())
+        logger.error(f"Reward function '{name}' not found. Available: {available}")
+    return func
+
+
+def get_verifier_function(name: str) -> Optional[VerifierFunction]:
+    """Retrieve a registered verifier function by name."""
+    func = VERIFIER_REGISTRY.get(name)
+    if func is None:
+        available = list(VERIFIER_REGISTRY.keys())
+        logger.error(f"Verifier function '{name}' not found. Available: {available}")
+    return func
+
+
+def apply_verifiers_to_reward(
+    original_reward_func: RewardFunction,
+    verifier_names: List[str],
+    penalty_on_failure: float = 0.0
+) -> RewardFunction:
+    """
+    Apply verifiers to a reward function following retrain's pre-filtering approach.
+    
+    If any verifier fails, returns the penalty value without executing the reward function.
+    This follows retrain's philosophy of efficient pre-filtering.
+    
+    Args:
+        original_reward_func: The reward function to wrap
+        verifier_names: List of verifier names to apply
+        penalty_on_failure: Value to return if any verifier fails
+        
+    Returns:
+        Wrapped reward function that applies verifiers first
+    """
+    # Load verifier functions
+    loaded_verifiers: List[VerifierFunction] = []
+    missing_verifiers = []
+    
+    for verifier_name in verifier_names:
+        verifier_func = get_verifier_function(verifier_name)
+        if verifier_func is None:
+            missing_verifiers.append(verifier_name)
+        else:
+            loaded_verifiers.append(verifier_func)
+    
+    if missing_verifiers:
+        available = list(VERIFIER_REGISTRY.keys())
+        raise ValueError(f"Verifiers {missing_verifiers} not found for reward '{original_reward_func.__name__}'. Available: {available}")
+    
+    @functools.wraps(original_reward_func)
+    def reward_with_verifiers(prompt: str, completion: str, example: Dict[str, Any], **kwargs) -> float:
+        # Apply verifiers first (pre-filtering approach)
+        for i, verifier_func in enumerate(loaded_verifiers):
+            verifier_name = verifier_names[i]
+            try:
+                if not verifier_func(prompt, completion, example):
+                    logger.debug(f"Verifier '{verifier_name}' failed for reward '{original_reward_func.__name__}'. Applying penalty: {penalty_on_failure}")
+                    return penalty_on_failure
+            except Exception as e:
+                logger.error(f"Verifier '{verifier_name}' errored: {e}. Applying penalty: {penalty_on_failure}")
+                return penalty_on_failure
+        
+        # All verifiers passed, execute reward function
+        try:
+            result = original_reward_func(prompt, completion, example, **kwargs)
+            return float(result)
+        except Exception as e:
+            logger.error(f"Reward function '{original_reward_func.__name__}' errored after verifiers passed: {e}. Returning 0.0")
+            return 0.0
+    
+    # Set descriptive name
+    if verifier_names:
+        verifier_suffix = '_and_'.join(verifier_names)
+        reward_with_verifiers.__name__ = f"{original_reward_func.__name__}_verified_by_{verifier_suffix}"
+    else:
+        reward_with_verifiers.__name__ = original_reward_func.__name__
+    
+    return reward_with_verifiers
+
+
+@dataclass
+class RewardConfig:
+    """Configuration for a reward function following retrain's patterns."""
+    name: str  # Name in REWARD_REGISTRY
+    weight: float = 1.0
+    params: Optional[Dict[str, Any]] = None
+    verifiers: Optional[List[str]] = None
+    verifier_penalty: float = 0.0
+    
+    def __post_init__(self):
+        if self.params is None:
+            self.params = {}
+        if self.verifiers is None:
+            self.verifiers = []
+
+
+def create_configured_reward_function(config: RewardConfig) -> RewardFunction:
+    """
+    Create a reward function from configuration following retrain's approach.
+    
+    This function:
+    1. Loads the base reward function from registry
+    2. Applies verifiers if specified
+    3. Handles parameter passing
+    4. Returns a configured function ready for use
+    """
+    # Get base reward function
+    base_reward_func = get_reward_function(config.name)
+    if base_reward_func is None:
+        raise ValueError(f"Reward function '{config.name}' not found in registry")
+    
+    # Create wrapper that handles params
+    def reward_with_params(prompt: str, completion: str, example: Dict[str, Any], **kwargs) -> float:
+        # Merge config params with runtime kwargs
+        merged_kwargs = {**config.params, **kwargs}
+        return base_reward_func(prompt, completion, example, **merged_kwargs)
+    
+    reward_with_params.__name__ = f"{base_reward_func.__name__}_with_params"
+    
+    # Apply verifiers if specified
+    if config.verifiers:
+        reward_with_params = apply_verifiers_to_reward(
+            reward_with_params, 
+            config.verifiers, 
+            config.verifier_penalty
+        )
+    
+    return reward_with_params
+
+
+# MLX optimization support
+@mx.compile
+def batch_reward_computation(
+    base_rewards: mx.array,
+    weights: mx.array
+) -> mx.array:
+    """
+    MLX-compiled function for efficient batch reward computation.
+    
+    Args:
+        base_rewards: Individual reward scores [batch_size, num_rewards]
+        weights: Reward weights [num_rewards]
+        
+    Returns:
+        Weighted combined rewards [batch_size]
+    """
+    return mx.sum(base_rewards * weights, axis=1)
+
+
+def list_registered_functions() -> Dict[str, List[str]]:
+    """List all registered reward and verifier functions."""
+    return {
+        "rewards": list(REWARD_REGISTRY.keys()),
+        "verifiers": list(VERIFIER_REGISTRY.keys())
+    }
+
+
+def clear_registries():
+    """Clear all registries (useful for testing)."""
+    global REWARD_REGISTRY, VERIFIER_REGISTRY
+    REWARD_REGISTRY.clear()
+    VERIFIER_REGISTRY.clear()
+    logger.info("Cleared all reward and verifier registries")

textpolicy/rewards/rollout_rewards.py

@@ -0,0 +1,410 @@
+# textpolicy/rewards/rollout_rewards.py
+"""
+Rollout-level reward processing system for efficient MLX training.
+
+This system processes rewards at the episode/rollout level rather than
+per-transition, enabling vectorized operations and batch processing
+for optimal MLX performance.
+
+Key features:
+- Batch reward computation for entire episodes
+- Vectorized operations using MLX
+- Integration with rollout buffer system
+- Support for async external reward models
+- Pure function composition
+"""
+
+from typing import Dict, List, Optional, Any
+import mlx.core as mx
+from dataclasses import dataclass
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+
+# Optional dependency
+try:
+    import aiohttp # type: ignore
+    HAS_AIOHTTP = True
+except ImportError:
+    HAS_AIOHTTP = False
+    aiohttp = None
+
+# Import reward functions used in this module
+# These are registered via the @reward decorator and provide the standard (prompt, completion, example) signature
+from .basic import length_reward, keyword_reward, perplexity_reward, accuracy_reward
+
+
+@dataclass
+class RewardConfig:
+    """Configuration for rollout reward processing."""
+    # Basic reward weights
+    length_weight: float = 0.1
+    keyword_weight: float = 0.2
+    perplexity_weight: float = 0.3
+    accuracy_weight: float = 0.4
+    
+    # Target parameters
+    target_length: int = 50
+    keywords: Optional[List[str]] = None
+    
+    # External reward model
+    external_rm_url: Optional[str] = None
+    external_rm_timeout: float = 30.0
+    
+    # Batch processing
+    batch_size: int = 32
+    max_workers: int = 4
+    
+    def __post_init__(self):
+        if self.keywords is None:
+            self.keywords = []
+
+
+class RolloutRewardProcessor:
+    """
+    Efficient rollout-level reward processor for MLX training.
+    
+    Processes entire episodes in batches, enabling vectorized operations
+    and optimal memory usage on Apple Silicon.
+    """
+    
+    def __init__(self, config: RewardConfig):
+        """
+        Initialize reward processor.
+        
+        Args:
+            config: Reward processing configuration
+        """
+        self.config = config
+        self.executor = ThreadPoolExecutor(max_workers=config.max_workers)
+        
+        # Pre-compile reward functions for MLX
+        self._compile_reward_functions()
+    
+    def _compile_reward_functions(self):
+        """Pre-compile reward functions for MLX optimization."""
+        # These will be compiled when first called
+        self._compiled_functions = {}
+    
+    def process_rollout(
+        self, 
+        prompt: str, 
+        completion: str, 
+        example: Dict[str, Any]
+    ) -> mx.array:
+        """
+        Process a single rollout with prompt, completion, and example context.
+        
+        This method provides the standard interface expected by tests and other components.
+        It follows the same signature as individual reward functions for consistency.
+        
+        Args:
+            prompt: Input prompt text
+            completion: Generated completion text
+            example: Example context with target parameters and metadata
+            
+        Returns:
+            MLX array with single reward value [1]
+        """
+        # Convert single rollout to episode format for consistency
+        episode = {
+            'prompt': prompt,
+            'response': completion,
+            'metadata': example
+        }
+        
+        # Process as single episode and return scalar reward
+        rewards = self.process_episode_rewards([episode])
+        return rewards[0] if rewards.size > 0 else mx.array(0.0)
+    
+    def process_batch_rollouts(
+        self, 
+        prompts: List[str], 
+        completions: List[str], 
+        examples: List[Dict[str, Any]]
+    ) -> mx.array:
+        """
+        Process multiple rollouts in batch for efficient processing.
+        
+        This method provides batch processing capability expected by tests and training loops.
+        It maintains the same interface as individual reward functions but processes multiple
+        examples simultaneously for better performance.
+        
+        Args:
+            prompts: List of input prompt texts
+            completions: List of generated completion texts
+            examples: List of example contexts with target parameters
+            
+        Returns:
+            MLX array of rewards [num_rollouts]
+        """
+        if len(prompts) != len(completions) or len(completions) != len(examples):
+            raise ValueError(f"Mismatched input lengths: prompts={len(prompts)}, "
+                           f"completions={len(completions)}, examples={len(examples)}")
+        
+        # Convert to episode format for consistency with existing implementation
+        episodes = []
+        for prompt, completion, example in zip(prompts, completions, examples):
+            episode = {
+                'prompt': prompt,
+                'response': completion,
+                'metadata': example
+            }
+            episodes.append(episode)
+        
+        # Process batch using existing episode processing logic
+        return self.process_episode_rewards(episodes)
+    
+    def process_episode_rewards(
+        self, 
+        episodes: List[Dict[str, Any]]
+    ) -> mx.array:
+        """
+        Process rewards for a batch of episodes.
+        
+        Args:
+            episodes: List of episode dictionaries with 'prompt' and 'response' fields
+            
+        Returns:
+            MLX array of rewards [num_episodes]
+        """
+        if not episodes:
+            return mx.array([])
+        
+        # Extract prompts and responses
+        prompts = [ep.get('prompt', '') for ep in episodes]
+        responses = [ep.get('response', '') for ep in episodes]
+        
+        # Compute basic rewards (vectorized where possible)
+        rewards = self._compute_basic_rewards(prompts, responses)
+        
+        # Add external reward model scores if configured
+        if self.config.external_rm_url:
+            external_rewards = self._get_external_rewards(episodes)
+            # Blend with basic rewards
+            alpha = 0.7  # Weight for external rewards
+            rewards = alpha * external_rewards + (1 - alpha) * rewards
+        
+        return rewards
+    
+    def _compute_basic_rewards(
+        self, 
+        prompts: List[str], 
+        responses: List[str]
+    ) -> mx.array:
+        """
+        Compute basic rewards using vectorized operations.
+        
+        Args:
+            prompts: List of input prompts
+            responses: List of generated responses
+            
+        Returns:
+            MLX array of combined rewards
+        """
+        # Initialize reward array
+        num_episodes = len(prompts)
+        rewards = mx.zeros(num_episodes)
+        
+        # Length rewards - pass config as example dict for unified interface
+        # The @reward decorator enforces (prompt, completion, example) signature
+        if self.config.length_weight > 0:
+            length_rewards = mx.array([
+                length_reward(p, r, {"target_length": self.config.target_length}) # type: ignore
+                for p, r in zip(prompts, responses)
+            ])
+            rewards += self.config.length_weight * length_rewards
+        
+        # Keyword rewards - pass config as example dict for unified interface  
+        # The @reward decorator enforces (prompt, completion, example) signature
+        if self.config.keyword_weight > 0 and self.config.keywords:
+            keyword_rewards = mx.array([
+                keyword_reward(p, r, {"keywords": self.config.keywords}) # type: ignore
+                for p, r in zip(prompts, responses)
+            ])
+            rewards += self.config.keyword_weight * keyword_rewards
+        
+        # Perplexity rewards - pass empty example dict for unified interface
+        # The @reward decorator enforces (prompt, completion, example) signature
+        if self.config.perplexity_weight > 0:
+            perplexity_rewards = mx.array([
+                perplexity_reward(p, r, {}) # type: ignore
+                for p, r in zip(prompts, responses)
+            ])
+            rewards += self.config.perplexity_weight * perplexity_rewards
+        
+        # Accuracy rewards - pass empty example dict for unified interface
+        # The @reward decorator enforces (prompt, completion, example) signature
+        if self.config.accuracy_weight > 0:
+            accuracy_rewards = mx.array([
+                accuracy_reward(p, r, {}) # type: ignore
+                for p, r in zip(prompts, responses)
+            ])
+            rewards += self.config.accuracy_weight * accuracy_rewards
+        
+        return rewards
+    
+    def _get_external_rewards(
+        self, 
+        episodes: List[Dict[str, Any]]
+    ) -> mx.array:
+        """
+        Get rewards from external reward model.
+        
+        Args:
+            episodes: List of episode dictionaries
+            
+        Returns:
+            MLX array of external rewards
+        """
+        # For now, return default rewards
+        # In practice, this would call external API or model
+        return mx.ones(len(episodes)) * 0.5
+    
+    async def _async_external_rewards(
+        self, 
+        episodes: List[Dict[str, Any]]
+    ) -> List[float]:
+        """
+        Async version for external reward model calls.
+        
+        Args:
+            episodes: List of episode dictionaries
+            
+        Returns:
+            List of reward scores
+        """
+        if not self.config.external_rm_url:
+            return [0.5] * len(episodes)
+        
+        if not HAS_AIOHTTP or aiohttp is None:
+            return [0.5] * len(episodes)
+        
+        # Type guard: aiohttp is guaranteed to be available from this point
+        # Assert for type checker that aiohttp is not None after the above checks
+        assert aiohttp is not None, "aiohttp should be available when HAS_AIOHTTP is True"
+        
+        async def get_reward(session, episode):
+            payload = {
+                "prompt": episode.get("prompt", ""),
+                "response": episode.get("response", ""),
+                "metadata": episode.get("metadata", {})
+            }
+            
+            try:
+                async with session.post(
+                    self.config.external_rm_url,
+                    json=payload,
+                    timeout=aiohttp.ClientTimeout(total=self.config.external_rm_timeout)
+                ) as resp:
+                    if resp.status == 200:
+                        result = await resp.json()
+                        return result.get("reward", 0.5)
+                    else:
+                        return 0.5
+            except Exception:
+                return 0.5
+        
+        async with aiohttp.ClientSession() as session:
+            tasks = [get_reward(session, ep) for ep in episodes]
+            results = await asyncio.gather(*tasks)
+            return list(results)
+    
+    def process_buffer_rewards(self, buffer) -> mx.array:
+        """
+        Process rewards for all episodes in a buffer.
+        
+        Args:
+            buffer: Buffer instance containing episodes
+            
+        Returns:
+            MLX array of episode rewards
+        """
+        # Extract episodes from buffer
+        episodes = []
+        for episode in buffer.storage.episodes:
+            # Convert episode to dict format
+            episode_dict = {
+                'prompt': episode.obs[0] if episode.obs else '',
+                'response': episode.act[-1] if episode.act else '',
+                'metadata': {
+                    'length': len(episode.obs),
+                    'logprobs': episode.logprob,
+                    'values': episode.value
+                }
+            }
+            episodes.append(episode_dict)
+        
+        return self.process_episode_rewards(episodes)
+    
+    def close(self):
+        """Cleanup resources."""
+        self.executor.shutdown(wait=True)
+
+
+# Pure function interface for integration with rollout system
+def create_rollout_reward_processor(config: RewardConfig) -> RolloutRewardProcessor:
+    """
+    Factory function for creating reward processors.
+    
+    Args:
+        config: Reward processing configuration
+        
+    Returns:
+        RolloutRewardProcessor instance
+    """
+    return RolloutRewardProcessor(config)
+
+
+def process_episode_batch_rewards(
+    episodes: List[Dict[str, Any]],
+    config: RewardConfig
+) -> mx.array:
+    """
+    Pure function for processing episode rewards.
+    
+    Args:
+        episodes: List of episode dictionaries
+        config: Reward configuration
+        
+    Returns:
+        MLX array of rewards
+    """
+    processor = RolloutRewardProcessor(config)
+    try:
+        return processor.process_episode_rewards(episodes)
+    finally:
+        processor.close()
+
+
+# MLX-compiled reward computation for high-performance training
+@mx.compile
+def compute_reward_vector(
+    response_lengths: mx.array,
+    keyword_matches: mx.array,
+    fluency_scores: mx.array,
+    accuracy_scores: mx.array,
+    weights: mx.array
+) -> mx.array:
+    """
+    MLX-compiled function for vectorized reward computation.
+    
+    Args:
+        response_lengths: Normalized length scores [batch_size]
+        keyword_matches: Keyword match scores [batch_size]
+        fluency_scores: Fluency scores [batch_size]
+        accuracy_scores: Accuracy scores [batch_size]
+        weights: Reward weights [4] (length, keyword, fluency, accuracy)
+        
+    Returns:
+        Combined rewards [batch_size]
+    """
+    # Weighted combination of reward components
+    rewards = (
+        weights[0] * response_lengths +
+        weights[1] * keyword_matches +
+        weights[2] * fluency_scores +
+        weights[3] * accuracy_scores
+    )
+    
+    return rewards
+