eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/reward_function.py

@@ -0,0 +1,410 @@
+import importlib
+import importlib.util
+import inspect
+import logging
+import os
+import warnings
+from functools import wraps
+from typing import Any, Callable, Dict, List, Optional, TypeVar, Union, cast
+
+import requests
+
+from .models import EvaluateResult, MetricResult
+from .typed_interface import reward_function
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=Callable[..., EvaluateResult])
+
+
+class RewardFunction:
+    """
+    A wrapper for reward functions that allows them to be run locally or remotely.
+
+    The RewardFunction class wraps a reward function (either a local function or a remote endpoint)
+    and provides a unified interface for calling it. It supports:
+
+    - Local functions (mode="local")
+    - Remote endpoints (mode="remote")
+    - Fireworks-hosted models (mode="fireworks_hosted")
+
+    Args:
+        func: The local function to use (for mode="local")
+        func_path: A string path to a function (e.g., "module.submodule:function_name")
+        mode: The mode of operation ("local", "remote", or "fireworks_hosted")
+        endpoint: The URL of the remote endpoint (for mode="remote")
+        model_id: The ID of the Fireworks-hosted model (for mode="fireworks_hosted")
+        **kwargs: Additional keyword arguments to pass to the function
+    """
+
+    def __init__(
+        self,
+        func: Optional[Callable] = None,
+        func_path: Optional[str] = None,
+        mode: str = "local",
+        endpoint: Optional[str] = None,
+        name: Optional[str] = None,
+        model_id: Optional[str] = None,
+        **kwargs,
+    ):
+        self.mode = mode
+        self.func = func
+        self.func_path = func_path
+        self.endpoint = endpoint
+        self.name = name
+        self.model_id = model_id
+        self.kwargs = kwargs
+
+        if mode == "local":
+            if func is None and func_path is None:
+                raise ValueError("Either 'func' or 'func_path' must be provided for local mode")
+            if func_path and func is None:
+                self.func = self._load_function_from_path(func_path)
+        elif mode == "remote":
+            if endpoint is None and name is None:
+                raise ValueError("Either 'endpoint' or 'name' must be provided for remote mode")
+            if name and endpoint is None:
+                self.endpoint = f"https://api.fireworks.ai/v1/reward/{name}"
+        elif mode == "fireworks_hosted":
+            if model_id is None:
+                raise ValueError("'model_id' must be provided for fireworks_hosted mode")
+            self.endpoint = f"https://api.fireworks.ai/v1/models/{model_id}/reward"
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+
+    def _load_function_from_path(self, func_path: str) -> Callable:
+        """
+        Load a function from a path string.
+        The path string should be in the format 'module.submodule:function_name' or 'module.submodule.function_name'.
+        """
+        # Check for the colon format first (preferred)
+        if ":" in func_path:
+            module_path, func_name = func_path.split(":", 1)
+
+            try:
+                module = importlib.import_module(module_path)
+                func = getattr(module, func_name)
+                return func
+            except (ImportError, AttributeError) as e:
+                raise ImportError(f"Failed to load function from path {func_path}: {str(e)}")
+
+        # Try dot notation format: module.path.function_name
+        # This assumes the last component is the function name
+        parts = func_path.split(".")
+        if len(parts) < 2:
+            raise ValueError(
+                f"Invalid func_path format: {func_path}, expected 'module.path:function_name' or 'module.path.function_name'"
+            )
+
+        module_path = ".".join(parts[:-1])
+        func_name = parts[-1]
+
+        try:
+            module = importlib.import_module(module_path)
+            func = getattr(module, func_name)
+            return func
+        except (ImportError, AttributeError) as e:
+            raise ImportError(f"Failed to load function from path {func_path}: {str(e)}")
+
+    def __call__(
+        self,
+        messages: List[Dict[str, str]],
+        ground_truth: Optional[Union[str, List[Dict[str, str]]]] = None,
+        **kwargs,
+    ) -> EvaluateResult:
+        """
+        Call the reward function with the provided messages.
+
+        Args:
+            messages: List of conversation messages, each with 'role' and 'content' keys
+            ground_truth: Ground truth data, which can be a string (e.g., an expected answer)
+                          or a list of original conversation messages (for context).
+                          If None and context is expected as a list, defaults to messages[:-1].
+            **kwargs: Additional keyword arguments to pass to the function
+
+        Returns:
+            EvaluateResult object with score and metrics
+        """
+        if ground_truth is None:
+            # Default to messages[:-1] if ground_truth is not provided and expected as context (list)
+            # This maintains previous behavior of original_messages defaulting.
+            # If ground_truth is meant to be a string and is None, it should be handled by the specific reward function.
+            ground_truth = messages[:-1] if messages else []
+
+        combined_kwargs = {**self.kwargs, **kwargs}
+
+        if self.mode == "local":
+            if self.func is None:
+                raise ValueError("No function provided for local mode")
+
+            try:
+                result = self.func(
+                    messages=messages,
+                    ground_truth=ground_truth,
+                    **combined_kwargs,
+                )
+
+                if isinstance(result, EvaluateResult):
+                    return result
+                elif isinstance(result, tuple) and len(result) == 2:
+                    # Handle legacy (score, components) tuple format
+                    warnings.warn(
+                        "Tuple return format is deprecated. Use EvaluateResult instead.",
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    score, components = result
+                    metrics = {
+                        k: MetricResult(score=v, reason=f"{k} score", is_score_valid=True)
+                        for k, v in components.items()
+                    }
+                    return EvaluateResult(score=score, metrics=metrics)
+                elif isinstance(result, dict) and "score" in result:
+                    # Handle dictionary return format
+                    warnings.warn(
+                        "Dictionary return format is deprecated. Use EvaluateResult instead.",
+                        DeprecationWarning,
+                        stacklevel=2,
+                    )
+                    metrics = {}
+                    if "metrics" in result:
+                        for k, v in result["metrics"].items():
+                            if isinstance(v, dict):
+                                metrics[k] = MetricResult(
+                                    score=v.get("score", 0.0),
+                                    reason=v.get("reason", f"{k} score"),
+                                    is_score_valid=v.get("is_score_valid", True),
+                                )
+                            else:
+                                metrics[k] = MetricResult(
+                                    score=float(v),
+                                    reason=f"{k} score",
+                                    is_score_valid=True,
+                                )
+                    return EvaluateResult(
+                        score=result["score"],
+                        reason=result.get("reason"),
+                        metrics=metrics,
+                    )
+                else:
+                    raise TypeError(
+                        f"Invalid return type from reward function: {type(result)}. "
+                        f"Expected EvaluateResult or (float, Dict[str, float]) tuple."
+                    )
+
+            except Exception as e:
+                logger.error(f"Error calling local reward function: {str(e)}")
+                raise
+
+        elif self.mode in ["remote", "fireworks_hosted"]:
+            if self.endpoint is None:
+                raise ValueError(f"No endpoint provided for {self.mode} mode")
+
+            payload = {
+                "messages": messages,
+                "ground_truth": ground_truth,
+                **combined_kwargs,
+            }
+
+            api_key = os.environ.get("FIREWORKS_API_KEY")
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}" if api_key else "",
+            }
+
+            try:
+                response = requests.post(self.endpoint, json=payload, headers=headers)
+                response.raise_for_status()
+                result = response.json()
+
+                if isinstance(result, dict) and "score" in result:
+                    metrics = {}
+                    if "metrics" in result:
+                        for k, v in result["metrics"].items():
+                            if isinstance(v, dict):
+                                metrics[k] = MetricResult(
+                                    score=v.get("score", 0.0),
+                                    reason=v.get("reason", f"{k} score"),
+                                    is_score_valid=v.get("success", True),
+                                )
+                            else:
+                                metrics[k] = MetricResult(
+                                    score=float(v),
+                                    reason=f"{k} score",
+                                    is_score_valid=True,
+                                )
+
+                    return EvaluateResult(
+                        score=result["score"],
+                        reason=result.get("reason"),
+                        metrics=metrics,
+                    )
+                else:
+                    raise ValueError(f"Invalid response from remote endpoint: {result}")
+
+            except Exception as e:
+                logger.error(f"Error calling remote endpoint: {str(e)}")
+                raise
+
+        raise ValueError(f"Invalid mode: {self.mode}")
+
+    def get_trl_adapter(self) -> Callable:
+        """
+        Create an adapter function for use with TRL library.
+
+        The TRL library expects a function that takes batch inputs and returns a batch of reward values.
+        This adapter handles:
+        1. Batch of messages (List[List[Dict]]) and original messages (List[List[Dict]])
+        2. Batch of texts (List[str]) for simpler cases
+
+        Returns:
+            A callable function compatible with TRL's expected signature for reward functions.
+        """
+
+        def adapter(prompts: List[List[Dict]], completions: Optional[List[str]] = None, **kwargs) -> List[float]:
+            """
+            Adapter function compatible with TRL's expected reward function signature.
+            TRL typically expects: (prompts: List[str], completions: List[str], **kwargs: Any) -> List[float]
+            This adapter handles the conversion from reward-kit's Message format.
+
+            Args:
+                prompts: A batch of prompt message lists as expected by this RewardFunction instance.
+                         Typically List[List[Dict[str, str]]], e.g.,
+                         [[{'role':'system',...}, {'role':'user',...}], ...]
+                completions: A batch of generated completion strings by the model.
+                             Optional. If None, it's assumed that the `prompts` argument
+                             already contains the full conversation history including the assistant's response.
+                **kwargs: Additional keyword arguments passed by TRL. These often include
+                          other columns from the HuggingFace dataset being used for training
+                          (e.g., 'solution', 'reference_answer'). These are expected to be
+                          lists of the same length as `prompts`.
+
+            Returns:
+                A list of float reward scores for the batch, one score per sample.
+            """
+            results = []
+            batch_size = len(prompts)
+
+            # If completions is None, assume prompts already contain complete conversations
+            if completions is None:
+                completions = [""] * batch_size
+
+            # Make sure completions has the right length after the None check
+            if batch_size != len(completions):
+                logger.warning(
+                    f"Batch size mismatch between prompts ({batch_size}) and "
+                    f"completions ({len(completions)}). Using min length."
+                )
+                batch_size = min(batch_size, len(completions))
+
+            # Extract potential ground truth solutions if available
+            # TRL passes columns from the dataset that weren't removed.
+            # We expect 'solution' based on our grpo_example.py setup.
+            solutions = kwargs.get("solution", [None] * batch_size)
+            if not isinstance(solutions, list) or len(solutions) != batch_size:
+                logger.warning(
+                    f"Expected 'solution' kwarg to be a list of size {batch_size}, but got {type(solutions)}. Ground truth might not be passed correctly."
+                )
+                solutions = [None] * batch_size
+
+            for i in range(batch_size):
+                completion_input = completions[i]
+                actual_completion_str = ""
+
+                if isinstance(completion_input, list):
+                    if completion_input:
+                        first_element = completion_input[0]
+                        if (
+                            isinstance(first_element, dict)
+                            and "content" in first_element
+                            and isinstance(first_element.get("role"), str)
+                            and first_element.get("role") == "assistant"
+                        ):
+                            # Expected structure: completions[i] = [{'role': 'assistant', 'content': 'str_content'}]
+                            actual_completion_str = str(first_element.get("content", ""))
+                            logger.debug(
+                                f"Adapter: completions[{i}] is a list with an assistant message dict. Extracted content."
+                            )
+                        else:
+                            logger.warning(
+                                f"Adapter: completions[{i}] is a list, but its first element "
+                                f"is not the expected assistant message dict or is malformed: {first_element}. "
+                                f"Using str(first_element) as content."
+                            )
+                            actual_completion_str = str(first_element)
+                    else:
+                        logger.warning(f"Adapter: completions[{i}] is an empty list. Using empty string for content.")
+                        actual_completion_str = ""
+                elif isinstance(completion_input, str):
+                    actual_completion_str = completion_input
+                else:
+                    # Fallback for other types (e.g. a direct dict, though less likely given warnings)
+                    logger.warning(
+                        f"Adapter: completions[{i}] is of unexpected type: {type(completion_input)}. "
+                        f"Attempting to stringify for content: {completion_input}"
+                    )
+                    actual_completion_str = str(completion_input)
+
+                messages = prompts[i] + [{"role": "assistant", "content": actual_completion_str}]
+
+                # Prepare kwargs for the underlying reward function call for this specific sample
+                call_kwargs: Dict[str, Any] = {}  # Initialize with Any type for values
+                current_solution = solutions[i]
+
+                debug_solution_val_str = str(current_solution) if current_solution is not None else "None"
+                logger.debug(
+                    f"Adapter loop i={i}, type(current_solution)={type(current_solution)}, value='{debug_solution_val_str[:100]}...'"
+                )
+
+                processed_solution_val: Optional[str] = None
+                if current_solution is not None:
+                    if isinstance(current_solution, list):
+                        logger.warning(
+                            f"Sample {i} solution is a list, attempting to use first element: {current_solution}"
+                        )
+                        if current_solution:
+                            processed_solution_val = str(current_solution[0])
+                        # If current_solution is an empty list, processed_solution_val remains None
+                    else:
+                        processed_solution_val = str(current_solution)
+
+                if processed_solution_val is not None:
+                    call_kwargs["solution"] = processed_solution_val
+
+                # Add any other necessary kwargs (if they were extracted from the main **kwargs)
+                # For now, only "solution" is dynamically handled from TRL kwargs.
+
+                try:
+                    # Call the underlying RewardFunction instance (__call__)
+                    # Pass the constructed messages and the extracted kwargs for this sample
+                    result = self(
+                        messages=messages,
+                        ground_truth=(
+                            call_kwargs.pop("solution") if "solution" in call_kwargs else None
+                        ),  # Pass solution as ground_truth if available
+                        **call_kwargs,
+                    )
+                    # Handle both RewardOutput and EvaluateResult
+                    score = result.score
+                    results.append(score)
+                except Exception as e:
+                    logger.error(f"Error processing sample {i} in TRL adapter: {str(e)}")
+                    # Append a default low score (e.g., 0.0) on error
+                    results.append(0.0)
+
+            return results
+
+        return adapter
+
+
+# The legacy_reward_function decorator has been removed as it is no longer needed.
+# Use the reward_function decorator from typed_interface instead.
+#
+# For deployment functionality, use the RewardFunction class or the deployment
+# methods from the evaluation module directly.
+
+
+# The alias below is removed to ensure that `from .typed_interface import reward_function`
+# is the one used throughout the library, thus avoiding the deprecation warning
+# when using the @reward_function decorator.

eval_protocol/rewards/__init__.py

@@ -0,0 +1,94 @@
+"""
+Out-of-the-box reward functions for common use cases.
+"""
+
+# Import specific reward functions
+from . import (
+    accuracy,
+    accuracy_length,
+    bfcl_reward,
+    code_execution,
+    cpp_code,
+    deepcoder_reward,
+    format,
+    function_calling,
+    json_schema,
+    language_consistency,
+    lean_prover,
+    length,
+    list_comparison_math_reward,
+    math,
+    multiple_choice_math_reward,
+    reasoning_steps,
+    repetition,
+    tag_count,
+)
+from .accuracy_length import cosine_scaled_accuracy_length_reward
+
+# Import function separately to avoid name conflict with the module
+from .bfcl_reward import bfcl_reward as bfcl_reward_function
+
+# Directly import specific reward functions for easy access
+from .code_execution import fractional_code_reward
+from .cpp_code import binary_cpp_code_reward, ioi_cpp_code_reward
+from .deepcoder_reward import deepcoder_code_reward
+
+# To make individual functions directly importable from eval_protocol.rewards
+# e.g., from eval_protocol.rewards import composite_function_call_reward
+from .function_calling import (
+    composite_function_call_reward,
+    exact_tool_match_reward,
+    llm_judge_reward,
+    schema_jaccard_reward,
+)
+from .lean_prover import (
+    deepseek_huggingface_prover_benchmark,
+    deepseek_prover_v2_reward,
+    lean_prover_reward,
+)
+
+# Import these with aliases to avoid name conflicts
+from .list_comparison_math_reward import (
+    list_comparison_math_reward as list_comparison_math_reward_function,
+)
+from .multiple_choice_math_reward import (
+    multiple_choice_math_reward as multiple_choice_math_reward_function,
+)
+
+__all__ = [
+    # Modules
+    "function_calling",
+    "json_schema",
+    "math",
+    "code_execution",
+    "format",
+    "tag_count",
+    "accuracy",
+    "language_consistency",
+    "reasoning_steps",
+    "length",
+    "repetition",
+    "cpp_code",
+    "accuracy_length",
+    "lean_prover",
+    "deepcoder_reward",
+    "multiple_choice_math_reward",
+    "list_comparison_math_reward",
+    "bfcl_reward",
+    # Specific functions for direct import
+    "fractional_code_reward",
+    "deepcoder_code_reward",
+    "multiple_choice_math_reward_function",
+    "list_comparison_math_reward_function",
+    "ioi_cpp_code_reward",
+    "binary_cpp_code_reward",
+    "cosine_scaled_accuracy_length_reward",
+    "lean_prover_reward",
+    "deepseek_prover_v2_reward",
+    "deepseek_huggingface_prover_benchmark",
+    "bfcl_reward_function",
+    "composite_function_call_reward",
+    "exact_tool_match_reward",
+    "llm_judge_reward",
+    "schema_jaccard_reward",
+]