eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/utils/batch_evaluation.py

@@ -0,0 +1,217 @@
+"""
+Utilities for running batch evaluation on transformed N-variant data.
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from ..models import EvaluateResult
+from ..utils.module_loader import load_function as load_reward_function
+
+logger = logging.getLogger(__name__)
+
+
+def run_batch_evaluation(
+    batch_jsonl_path: str,
+    reward_function_path: str,
+    output_path: str,
+    reward_function_kwargs: Optional[Dict[str, Any]] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Run batch evaluation on transformed N-variant data.
+
+    Args:
+        batch_jsonl_path: Path to the batch evaluation JSONL file
+        reward_function_path: Path to the batch reward function (e.g., "module.function")
+        output_path: Path to write the batch evaluation results
+        reward_function_kwargs: Additional kwargs for the reward function
+
+    Returns:
+        List of batch evaluation results
+    """
+    if reward_function_kwargs is None:
+        reward_function_kwargs = {}
+
+    # Load the batch reward function
+    reward_function = load_reward_function(reward_function_path)
+
+    # Verify it's a batch mode function
+    if not hasattr(reward_function, "_reward_function_mode"):
+        logger.warning(f"Reward function {reward_function_path} doesn't have mode metadata. Assuming batch mode.")
+    elif getattr(reward_function, "_reward_function_mode") != "batch":
+        raise ValueError(
+            f"Reward function {reward_function_path} is not configured for batch mode. Expected mode='batch'."
+        )
+
+    results = []
+
+    try:
+        with open(batch_jsonl_path, "r", encoding="utf-8") as f:
+            for line_num, line in enumerate(f, 1):
+                try:
+                    data = json.loads(line.strip())
+
+                    # Extract required fields
+                    request_id = data.get("request_id")
+                    rollouts_messages = data.get("rollouts_messages")
+
+                    if not request_id:
+                        logger.error(f"Line {line_num}: Missing request_id")
+                        continue
+
+                    if not rollouts_messages or not isinstance(rollouts_messages, list):
+                        logger.error(f"Line {line_num}: Missing or invalid rollouts_messages")
+                        continue
+
+                    # Prepare kwargs for the batch reward function
+                    batch_kwargs = dict(reward_function_kwargs)
+
+                    # Add other fields from the data as kwargs (excluding the main inputs)
+                    excluded_fields = {
+                        "request_id",
+                        "rollouts_messages",
+                        "num_variants",
+                        "response_ids",
+                    }
+                    for key, value in data.items():
+                        if key not in excluded_fields:
+                            batch_kwargs[key] = value
+
+                    # Call the batch reward function
+                    try:
+                        batch_results = reward_function(rollouts_messages=rollouts_messages, **batch_kwargs)
+
+                        # Validate results
+                        if not isinstance(batch_results, list):
+                            raise ValueError(f"Batch reward function must return a list, got {type(batch_results)}")
+
+                        if len(batch_results) != len(rollouts_messages):
+                            raise ValueError(
+                                f"Batch reward function returned {len(batch_results)} results "
+                                f"but expected {len(rollouts_messages)} (one per rollout)"
+                            )
+
+                        # Create result entries
+                        response_ids = data.get("response_ids", list(range(len(rollouts_messages))))
+
+                        for i, (response_id, eval_result) in enumerate(zip(response_ids, batch_results)):
+                            if not isinstance(eval_result, EvaluateResult):
+                                logger.error(f"Result {i} is not an EvaluateResult: {type(eval_result)}")
+                                continue
+
+                            result_entry = {
+                                "request_id": request_id,
+                                "response_id": response_id,
+                                "rollout_index": i,
+                                "evaluation_score": eval_result.score,
+                                "evaluation_reason": eval_result.reason,
+                                "is_score_valid": eval_result.is_score_valid,
+                                "evaluation_metrics": (
+                                    {k: v.model_dump() for k, v in eval_result.metrics.items()}
+                                    if eval_result.metrics
+                                    else {}
+                                ),
+                                # Include original metadata
+                                **{k: v for k, v in data.items() if k not in excluded_fields},
+                            }
+
+                            results.append(result_entry)
+
+                    except Exception as e:
+                        logger.error(f"Error calling batch reward function for request {request_id}: {e}")
+                        # Create error entries for each expected result
+                        response_ids = data.get("response_ids", list(range(len(rollouts_messages))))
+                        for i, response_id in enumerate(response_ids):
+                            error_entry = {
+                                "request_id": request_id,
+                                "response_id": response_id,
+                                "rollout_index": i,
+                                "error": f"Batch evaluation failed: {str(e)}",
+                                "evaluation_score": 0.0,
+                                "is_score_valid": False,
+                            }
+                            results.append(error_entry)
+
+                except json.JSONDecodeError as e:
+                    logger.error(f"Invalid JSON on line {line_num}: {e}")
+                    continue
+                except Exception as e:
+                    logger.error(f"Error processing line {line_num}: {e}")
+                    continue
+
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Batch JSONL file not found: {batch_jsonl_path}")
+
+    # Write results
+    with open(output_path, "w", encoding="utf-8") as f:
+        for result in results:
+            f.write(json.dumps(result) + "\n")
+
+    logger.info(f"Batch evaluation completed. {len(results)} results written to {output_path}")
+    return results
+
+
+def create_sample_batch_reward_function():
+    """
+    Create a sample batch reward function for testing.
+    This is a simple function that compares all variants and scores them relative to each other.
+    """
+    from ..models import EvaluateResult, Message
+    from ..typed_interface import reward_function
+
+    @reward_function(mode="batch")
+    def sample_batch_reward(
+        rollouts_messages: List[List[Message]],
+        ground_truth_for_eval: Optional[str] = None,
+        **kwargs: Any,
+    ) -> List[EvaluateResult]:
+        """
+        Sample batch reward function that scores variants relative to each other.
+
+        This function demonstrates how to process multiple rollouts (variants) together
+        and return comparative scores.
+        """
+        from ..models import MetricResult
+
+        results = []
+
+        # Extract the assistant responses from each rollout
+        assistant_responses = []
+        for rollout in rollouts_messages:
+            assistant_msg = None
+            for msg in reversed(rollout):  # Find the last assistant message
+                if msg.role == "assistant":
+                    assistant_msg = msg.content
+                    break
+            assistant_responses.append(assistant_msg or "")
+
+        # Simple scoring: longer responses get higher scores (just for demonstration)
+        response_lengths = [len(response) for response in assistant_responses]
+        max_length = max(response_lengths) if response_lengths else 1
+
+        for i, (response, length) in enumerate(zip(assistant_responses, response_lengths)):
+            # Normalize score based on length (0.1 to 1.0)
+            base_score = 0.1 + 0.9 * (length / max_length)
+
+            # Add some variation based on position (earlier variants get slight bonus)
+            position_bonus = 0.1 * (1 - i / len(assistant_responses))
+            final_score = min(1.0, base_score + position_bonus)
+
+            result = EvaluateResult(
+                score=final_score,
+                reason=f"Variant {i}: Length={length}, Base={base_score:.2f}, Position bonus={position_bonus:.2f}",
+                is_score_valid=True,
+                metrics={
+                    "response_length": MetricResult(
+                        score=min(1.0, length / 100.0),  # Normalize length score
+                        reason=f"Response contains {length} characters",
+                        is_score_valid=True,
+                    )
+                },
+            )
+            results.append(result)
+
+        return results
+
+    return sample_batch_reward

eval_protocol/utils/batch_transformation.py

@@ -0,0 +1,205 @@
+"""
+Utilities for transforming N-variant generation results into batch evaluation format.
+"""
+
+import json
+import logging
+from collections import defaultdict
+from typing import Any, Dict, List, Optional, Union
+
+logger = logging.getLogger(__name__)
+
+
+def transform_n_variant_jsonl_to_batch_format(
+    input_file_path: str,
+    output_file_path: Optional[str] = None,
+    request_id_field: str = "request_id",
+    response_id_field: str = "response_id",
+    messages_field: str = "full_conversation_history",
+    fallback_messages_fields: List[str] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Transform N-variant generation JSONL output into batch evaluation format.
+
+    This function groups variants by request_id and creates rollouts_messages
+    containing all variant conversations for each original request.
+
+    Args:
+        input_file_path: Path to the N-variant generation JSONL file
+        output_file_path: Optional path to write the transformed data (if None, returns data only)
+        request_id_field: Field name containing the original request ID (default: "request_id")
+        response_id_field: Field name containing the variant response ID (default: "response_id")
+        messages_field: Primary field containing conversation messages (default: "full_conversation_history")
+        fallback_messages_fields: Fallback fields to construct messages if primary field is missing
+
+    Returns:
+        List of batch evaluation entries, each containing rollouts_messages and other metadata
+
+    Raises:
+        FileNotFoundError: If input file doesn't exist
+        ValueError: If required fields are missing or data format is invalid
+    """
+    if fallback_messages_fields is None:
+        fallback_messages_fields = ["user_query", "system_prompt", "assistant_response"]
+
+    # Group variants by request_id
+    grouped_variants = defaultdict(list)
+
+    try:
+        with open(input_file_path, "r", encoding="utf-8") as f:
+            for line_num, line in enumerate(f, 1):
+                try:
+                    data = json.loads(line.strip())
+
+                    # Skip lines with errors
+                    if "error" in data:
+                        logger.warning(f"Skipping line {line_num} due to error: {data.get('error')}")
+                        continue
+
+                    # Validate required fields
+                    if request_id_field not in data:
+                        raise ValueError(f"Line {line_num}: Missing required field '{request_id_field}'")
+
+                    if response_id_field not in data:
+                        raise ValueError(f"Line {line_num}: Missing required field '{response_id_field}'")
+
+                    request_id = data[request_id_field]
+                    response_id = data[response_id_field]
+
+                    # Extract messages
+                    messages = _extract_messages_from_data(data, messages_field, fallback_messages_fields, line_num)
+
+                    # Store variant data
+                    variant_data = {
+                        "response_id": response_id,
+                        "messages": messages,
+                        "original_data": data,  # Keep original data for metadata extraction
+                    }
+
+                    grouped_variants[request_id].append(variant_data)
+
+                except json.JSONDecodeError as e:
+                    logger.error(f"Invalid JSON on line {line_num}: {e}")
+                    continue
+                except Exception as e:
+                    logger.error(f"Error processing line {line_num}: {e}")
+                    continue
+
+    except FileNotFoundError:
+        raise FileNotFoundError(f"Input file not found: {input_file_path}")
+
+    # Transform grouped variants into batch format
+    batch_entries = []
+
+    for request_id, variants in grouped_variants.items():
+        # Sort variants by response_id to ensure consistent ordering
+        variants.sort(key=lambda x: x["response_id"])
+
+        # Extract rollouts_messages (list of message lists)
+        rollouts_messages = [variant["messages"] for variant in variants]
+
+        # Extract common metadata from the first variant (assuming it's consistent across variants)
+        first_variant = variants[0]["original_data"]
+
+        # Create batch entry
+        batch_entry = {
+            "request_id": request_id,
+            "rollouts_messages": rollouts_messages,
+            "num_variants": len(variants),
+            "response_ids": [variant["response_id"] for variant in variants],
+        }
+
+        # Add common fields as kwargs (excluding variant-specific fields)
+        excluded_fields = {
+            "id",
+            request_id_field,
+            response_id_field,
+            messages_field,
+            "full_conversation_history",
+            "assistant_response",
+            "evaluation_score",
+            "evaluation_reason",
+            "evaluation_metrics",
+            "executed_tool_calls",
+            "discovered_tools",
+            "final_mcp_state_captured",
+        }
+
+        for key, value in first_variant.items():
+            if key not in excluded_fields:
+                batch_entry[key] = value
+
+        batch_entries.append(batch_entry)
+
+    # Write output file if specified
+    if output_file_path:
+        with open(output_file_path, "w", encoding="utf-8") as f:
+            for entry in batch_entries:
+                f.write(json.dumps(entry) + "\n")
+        logger.info(f"Transformed {len(batch_entries)} batch entries written to {output_file_path}")
+
+    return batch_entries
+
+
+def _extract_messages_from_data(
+    data: Dict[str, Any], primary_field: str, fallback_fields: List[str], line_num: int
+) -> List[Dict[str, Any]]:
+    """
+    Extract conversation messages from variant data.
+
+    Args:
+        data: Variant data dictionary
+        primary_field: Primary field containing messages
+        fallback_fields: Fallback fields to construct messages
+        line_num: Line number for error reporting
+
+    Returns:
+        List of message dictionaries
+    """
+    # Try primary field first
+    if primary_field in data and data[primary_field]:
+        messages = data[primary_field]
+        if isinstance(messages, list):
+            return messages
+        else:
+            logger.warning(f"Line {line_num}: {primary_field} is not a list, trying fallback")
+
+    # Try to construct messages from fallback fields
+    messages = []
+
+    # Add system message if available
+    if "system_prompt" in data and data["system_prompt"]:
+        messages.append({"role": "system", "content": data["system_prompt"]})
+
+    # Add user message if available
+    if "user_query" in data and data["user_query"]:
+        messages.append({"role": "user", "content": data["user_query"]})
+
+    # Add assistant message if available
+    if "assistant_response" in data and data["assistant_response"]:
+        messages.append({"role": "assistant", "content": data["assistant_response"]})
+
+    if not messages:
+        raise ValueError(f"Line {line_num}: Could not extract messages from any available fields")
+
+    return messages
+
+
+def create_batch_evaluation_dataset(n_variant_jsonl_path: str, output_jsonl_path: str, **transform_kwargs) -> str:
+    """
+    Convenience function to create a batch evaluation dataset from N-variant generation output.
+
+    Args:
+        n_variant_jsonl_path: Path to N-variant generation JSONL file
+        output_jsonl_path: Path for the batch evaluation JSONL file
+        **transform_kwargs: Additional arguments for transform_n_variant_jsonl_to_batch_format
+
+    Returns:
+        Path to the created batch evaluation dataset
+    """
+    transform_n_variant_jsonl_to_batch_format(
+        input_file_path=n_variant_jsonl_path,
+        output_file_path=output_jsonl_path,
+        **transform_kwargs,
+    )
+    return output_jsonl_path

eval_protocol/utils/dataset_helpers.py

@@ -0,0 +1,112 @@
+import json
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
+
+if TYPE_CHECKING:
+    from datasets import Dataset
+
+try:
+    from datasets import Dataset
+
+    HAS_DATASETS_LIB = True
+except ImportError:
+    HAS_DATASETS_LIB = False
+    if not TYPE_CHECKING:
+
+        class Dataset:
+            """Placeholder for HuggingFace Dataset for when the library is not installed."""
+
+            def __init__(self, *args: Any, **kwargs: Any) -> None:
+                pass
+
+
+def load_jsonl_to_hf_dataset(
+    dataset_path: str,
+    transform_fn: Optional[Callable[[Dict[str, Any]], Optional[Dict[str, Any]]]] = None,
+    prompt_column: str = "prompt",
+    required_columns: Optional[List[str]] = None,
+    dataset_filter_fn: Optional[Callable[[Dict[str, Any]], bool]] = None,
+) -> Optional["Dataset"]:
+    """
+    Loads a JSONL file into a HuggingFace Dataset, optionally applying a
+    transformation to each sample and ensuring required columns are present.
+
+    Args:
+        dataset_path: Path to the JSONL file.
+        transform_fn: An optional function to apply to each raw dictionary (sample)
+                      from the JSONL file. It should take a dict and return a dict
+                      (or None to skip the sample).
+        prompt_column: The name of the column expected to contain the prompt for TRL.
+                       Defaults to "prompt".
+        required_columns: A list of column names that must be present in the
+                          final dataset (after transformation).
+        dataset_filter_fn: An optional function to filter samples from the dataset
+                           after transformation. It should take a dict and return
+                           True to keep the sample, False to discard.
+
+    Returns:
+        A HuggingFace Dataset object, or None if the datasets library is not installed
+        or if an error occurs.
+    """
+    if not HAS_DATASETS_LIB:
+        print("The 'datasets' library is not installed. Please install it with 'pip install datasets'.")
+        return None
+
+    processed_samples: List[Dict[str, Any]] = []
+    try:
+        with open(dataset_path, "r", encoding="utf-8") as f:
+            for line_number, line in enumerate(f, 1):
+                try:
+                    raw_sample = json.loads(line.strip())
+
+                    transformed_sample: Optional[Dict[str, Any]]
+                    if transform_fn:
+                        transformed_sample = transform_fn(raw_sample)
+                    else:
+                        transformed_sample = raw_sample
+
+                    if transformed_sample is None:
+                        continue
+
+                    if dataset_filter_fn and not dataset_filter_fn(transformed_sample):
+                        continue
+
+                    processed_samples.append(transformed_sample)
+
+                except json.JSONDecodeError as e:
+                    print(f"Warning: Skipping line {line_number} in {dataset_path} due to JSON decode error: {e}")
+                except Exception as e:
+                    print(
+                        f"Warning: Skipping line {line_number} in {dataset_path} due to error in transform_fn or filter_fn: {e}"
+                    )
+
+        if not processed_samples:
+            print(f"Warning: No samples were processed from {dataset_path}.")
+            return Dataset.from_list([])
+
+        hf_dataset = Dataset.from_list(processed_samples)
+
+        if prompt_column not in hf_dataset.column_names:
+            raise ValueError(
+                f"Dataset from {dataset_path} must contain a '{prompt_column}' column after transformation."
+            )
+
+        final_required_columns = set(required_columns or [])
+        final_required_columns.add(prompt_column)
+
+        for col in final_required_columns:
+            if col not in hf_dataset.column_names:
+                raise ValueError(
+                    f"Dataset from {dataset_path} must contain a '{col}' column after transformation for the reward function/TRL."
+                )
+
+        return hf_dataset
+
+    except FileNotFoundError:
+        print(f"Error: Dataset file not found at {dataset_path}")
+        return None
+    except ValueError as ve:
+        print(f"Error: {ve}")
+        return None
+    except Exception as e:
+        print(f"An unexpected error occurred while loading dataset {dataset_path}: {e}")
+        return None

eval_protocol/utils/module_loader.py

@@ -0,0 +1,56 @@
+"""
+Utility for dynamically loading modules and functions.
+"""
+
+import importlib
+import logging
+from typing import Any, Callable
+
+logger = logging.getLogger(__name__)
+
+
+def load_function(import_path: str) -> Callable[..., Any]:
+    """
+    Dynamically loads a function given its full import path.
+    Example: "my_package.my_module.my_function"
+    """
+    try:
+        module_path, function_name = import_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        func = getattr(module, function_name)
+        if not callable(func):
+            raise AttributeError(f"'{function_name}' in module '{module_path}' is not callable.")
+        logger.info(f"Successfully loaded function '{function_name}' from '{module_path}'.")
+        return func
+    except ImportError as e:
+        logger.error(f"Failed to import module from path '{import_path}': {e}")
+        raise
+    except AttributeError as e:
+        logger.error(f"Failed to find or access function in path '{import_path}': {e}")
+        raise
+    except Exception as e:
+        logger.error(f"An unexpected error occurred while loading function from '{import_path}': {e}")
+        raise
+
+
+# Example usage:
+# if __name__ == '__main__':
+#     try:
+#         # Assuming you have a eval_protocol.rewards.math module with math_reward function
+#         math_reward_func = load_function("eval_protocol.rewards.math.math_reward")
+#         print(f"Loaded: {math_reward_func}")
+#         # You could then call it, e.g., if it took simple args: math_reward_func(arg1="test")
+#     except Exception as e:
+#         print(f"Test loading failed: {e}")
+
+#     try:
+#         # Test with a non-existent function
+#         load_function("eval_protocol.rewards.math.non_existent_function")
+#     except Exception as e:
+#         print(f"Test loading non-existent function failed as expected: {e}")
+
+#     try:
+#         # Test with a non-existent module
+#         load_function("non_existent_package.non_existent_module.some_function")
+#     except Exception as e:
+#         print(f"Test loading non-existent module failed as expected: {e}")