eval-protocol 0.0.3__py3-none-any.whl

eval_protocol/rewards/repetition.py

@@ -0,0 +1,342 @@
+"""
+Reward functions for evaluating repetition in model responses.
+
+This module provides reward functions that penalize repetitive text in model responses,
+encouraging more diverse and information-rich outputs.
+"""
+
+import re
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+
+from ..models import EvaluateResult, Message, MetricResult
+from ..typed_interface import reward_function
+
+
+def get_ngrams(text: str, n: int, language: str = "en") -> Tuple[List[Tuple[str, ...]], int]:
+    """
+    Extract n-grams from text based on language.
+
+    Args:
+        text: The text to extract n-grams from
+        n: Size of the n-grams
+        language: Language of the text (affects tokenization)
+
+    Returns:
+        Tuple of (list of n-grams, total n-gram count)
+    """
+    if language == "en":
+        words = text.lower().split()
+    elif language == "zh":
+        try:
+            import jieba
+
+            words = list(jieba.cut(text))
+        except ImportError:
+            words = list(text)
+    else:
+        words = text.lower().split()
+
+    ngrams = []
+    for i in range(len(words) - n + 1):
+        ngrams.append(tuple(words[i : i + n]))
+
+    return ngrams, len(ngrams)
+
+
+@reward_function
+def repetition_penalty_reward(
+    messages: Union[List[Message], List[Dict[str, Any]]],
+    ground_truth: Optional[
+        Union[List[Message], List[Dict[str, Any]]]
+    ] = None,  # Not used by this function but part of standard signature
+    ngram_size: int = 3,
+    max_penalty: float = 0.5,
+    language: str = "en",
+    **kwargs: Any,
+) -> EvaluateResult:
+    """
+    Reward function that penalizes repetitive text in model responses.
+    The model's response is assumed to be the last message in the `messages` list.
+
+    This function computes repetition by examining unique n-grams in the response
+    and penalizes texts with a high proportion of repeated phrases.
+
+    Args:
+        messages: List of conversation messages, where `messages[-1]` is the model's response.
+        ground_truth: Optional. Expected assistant response trajectory. Not directly used by this reward.
+        ngram_size: Size of n-grams to check for repetition.
+        max_penalty: Maximum penalty to apply for repetitive text.
+        language: Language of the text (affects tokenization).
+        **kwargs: Additional arguments.
+
+    Returns:
+        EvaluateResult with score penalizing repetition
+    """
+    if not messages or len(messages) == 0:
+        return EvaluateResult(
+            score=0.0,
+            reason="No messages provided",
+            metrics={"repetition": MetricResult(score=0.0, is_score_valid=False, reason="No messages provided")},
+        )
+
+    response = messages[-1]
+
+    if isinstance(response, Message):
+        if response.role != "assistant":
+            return EvaluateResult(
+                score=0.0,
+                reason="No assistant response found",
+                metrics={
+                    "repetition": MetricResult(
+                        score=0.0,
+                        is_score_valid=False,
+                        reason="Message not from assistant",
+                    )
+                },
+            )
+        text = response.content or ""
+    elif isinstance(response, dict):
+        if response.get("role") != "assistant":
+            return EvaluateResult(
+                score=0.0,
+                reason="No assistant response found",
+                metrics={
+                    "repetition": MetricResult(
+                        score=0.0,
+                        is_score_valid=False,
+                        reason="Message not from assistant",
+                    )
+                },
+            )
+        text = response.get("content", "")
+    else:
+        return EvaluateResult(
+            score=0.0,
+            reason="Last message is of unexpected type.",
+            metrics={
+                "repetition": MetricResult(
+                    score=0.0,
+                    is_score_valid=False,
+                    reason="Invalid message type in messages.",
+                )
+            },
+        )
+
+    if not text.strip():
+        return EvaluateResult(
+            score=1.0,
+            reason="Empty response, no repetition to penalize",
+            metrics={
+                "repetition": MetricResult(
+                    score=1.0,
+                    is_score_valid=True,
+                    reason="Empty response",
+                ),
+                "unique_ngram_ratio": MetricResult(
+                    score=1.0,
+                    is_score_valid=True,
+                    reason="Empty response",
+                ),
+                "repetition_penalty": MetricResult(
+                    score=1.0,
+                    is_score_valid=True,
+                    reason="No penalty applied to empty response",
+                ),
+            },
+        )
+
+    ngrams, total = get_ngrams(text, ngram_size, language)
+
+    if total < 1:
+        return EvaluateResult(
+            score=1.0,
+            reason=f"Text too short for {ngram_size}-gram analysis",
+            metrics={
+                "repetition": MetricResult(
+                    score=1.0,
+                    is_score_valid=True,
+                    reason=f"Text too short for {ngram_size}-gram analysis",
+                )
+            },
+        )
+
+    unique_ngrams = len(set(ngrams))
+    repetition_ratio = 1.0 - (unique_ngrams / total)
+    penalty = repetition_ratio * max_penalty
+    score = max(0.0, 1.0 - penalty)
+    success = repetition_ratio < 0.2
+
+    reason = f"Repetition ratio: {repetition_ratio:.2f}, Unique {ngram_size}-grams: {unique_ngrams}/{total}"
+    metrics = {
+        "repetition": MetricResult(score=score, is_score_valid=success, reason=reason),
+        "unique_ngram_ratio": MetricResult(
+            score=1.0 - repetition_ratio,
+            is_score_valid=success,
+            reason=f"Unique {ngram_size}-gram ratio: {1.0 - repetition_ratio:.2f}",
+        ),
+        "repetition_penalty": MetricResult(
+            score=1.0 - penalty,
+            is_score_valid=success,
+            reason=f"Applied repetition penalty: {penalty:.2f}",
+        ),
+    }
+
+    return EvaluateResult(score=score, reason=reason, metrics=metrics, is_score_valid=score > 0.0)
+
+
+@reward_function
+def diversity_reward(
+    messages: Union[List[Message], List[Dict[str, Any]]],
+    ground_truth: Optional[
+        Union[List[Message], List[Dict[str, Any]]]
+    ] = None,  # Not used by this function but part of standard signature
+    ngram_sizes: List[int] = [1, 2, 3],
+    weights: Optional[List[float]] = None,
+    language: str = "en",
+    **kwargs: Any,
+) -> EvaluateResult:
+    """
+    Reward function that measures lexical diversity in model responses.
+    The model's response is assumed to be the last message in the `messages` list.
+
+    This function computes diversity across multiple n-gram sizes and combines them
+    into a weighted score to encourage varied vocabulary and phrasing.
+
+    Args:
+        messages: List of conversation messages, where `messages[-1]` is the model's response.
+        ground_truth: Optional. Expected assistant response trajectory. Not directly used by this reward.
+        ngram_sizes: List of n-gram sizes to evaluate.
+        weights: Optional list of weights for each n-gram size (normalized if provided).
+        language: Language of the text (affects tokenization).
+        **kwargs: Additional arguments.
+
+    Returns:
+        EvaluateResult with score based on lexical diversity
+    """
+    if not messages or len(messages) == 0:
+        return EvaluateResult(
+            score=0.0,
+            reason="No messages provided",
+            metrics={"diversity": MetricResult(score=0.0, is_score_valid=False, reason="No messages provided")},
+        )
+
+    response = messages[-1]
+
+    if isinstance(response, Message):
+        if response.role != "assistant":
+            return EvaluateResult(
+                score=0.0,
+                reason="No assistant response found",
+                metrics={
+                    "diversity": MetricResult(
+                        score=0.0,
+                        is_score_valid=False,
+                        reason="Message not from assistant",
+                    )
+                },
+            )
+        text = response.content or ""
+    elif isinstance(response, dict):
+        if response.get("role") != "assistant":
+            return EvaluateResult(
+                score=0.0,
+                reason="No assistant response found",
+                metrics={
+                    "diversity": MetricResult(
+                        score=0.0,
+                        is_score_valid=False,
+                        reason="Message not from assistant",
+                    )
+                },
+            )
+        text = response.get("content", "")
+    else:
+        return EvaluateResult(
+            score=0.0,
+            reason="Last message is of unexpected type.",
+            metrics={
+                "diversity": MetricResult(
+                    score=0.0,
+                    is_score_valid=False,
+                    reason="Invalid message type in messages.",
+                )
+            },
+        )
+
+    if not text.strip():
+        return EvaluateResult(
+            score=0.0,
+            reason="Empty response",
+            metrics={
+                "diversity": MetricResult(
+                    score=0.0,
+                    is_score_valid=False,
+                    reason="Empty response",
+                )
+            },
+        )
+
+    if weights is None:
+        weights = [0.2, 0.3, 0.5][: len(ngram_sizes)]
+
+    if len(weights) != len(ngram_sizes):
+        if len(weights) > len(ngram_sizes):
+            weights = weights[: len(ngram_sizes)]
+        else:
+            missing_weight = (
+                (1.0 - sum(weights)) / (len(ngram_sizes) - len(weights))
+                if (len(ngram_sizes) - len(weights)) > 0
+                else 0
+            )
+            weights.extend([missing_weight] * (len(ngram_sizes) - len(weights)))
+
+    total_weight = sum(weights)
+    if total_weight != 1.0 and total_weight > 0:  # Avoid division by zero if total_weight is 0
+        weights = [w / total_weight for w in weights]
+    elif total_weight == 0 and len(weights) > 0:  # If all weights are zero, distribute equally
+        weights = [1.0 / len(weights)] * len(weights)
+
+    diversity_scores = {}
+    ratios = {}
+
+    for size, weight in zip(ngram_sizes, weights):
+        ngrams, total = get_ngrams(text, size, language)
+
+        if total < 1:
+            diversity_scores[f"ngram_{size}"] = 1.0
+            ratios[f"ngram_{size}"] = 1.0
+            continue
+
+        unique_count = len(set(ngrams))
+        ratio = unique_count / total
+        diversity_scores[f"ngram_{size}"] = ratio * weight
+        ratios[f"ngram_{size}"] = ratio
+
+    final_score = sum(diversity_scores.values())
+    success = final_score > 0.6
+
+    size_metric_items: List[Tuple[str, MetricResult]] = []
+    for size_key, ratio_val in ratios.items():
+        metric_for_size = MetricResult(
+            score=ratio_val,
+            is_score_valid=ratio_val > 0.7,
+            reason=f"Diversity ratio for {size_key}: {ratio_val:.2f}",
+        )
+        size_metric_items.append((size_key, metric_for_size))
+
+    size_metrics: Dict[str, MetricResult] = dict(size_metric_items)
+
+    metrics: Dict[str, MetricResult] = {
+        "diversity": MetricResult(
+            score=final_score,
+            is_score_valid=success,
+            reason=f"Overall weighted diversity score: {final_score:.2f}",
+        ),
+        **size_metrics,
+    }
+
+    return EvaluateResult(
+        score=final_score,
+        reason=f"Lexical diversity score: {final_score:.2f}",
+        metrics=metrics,
+    )

eval_protocol/rewards/tag_count.py

@@ -0,0 +1,162 @@
+"""
+Reward functions for counting tags in text responses.
+
+This module provides reward functions that evaluate if responses contain
+specified XML/HTML-like tags in correct quantities.
+"""
+
+import re
+from typing import Any, Dict, List, Set, Union
+
+from ..models import EvaluateResult, Message, MetricResult
+from ..typed_interface import reward_function
+
+
+@reward_function  # type: ignore[arg-type]
+def tag_count_reward(
+    messages: List[Message],
+    *,  # Make subsequent parameters keyword-only
+    required_tags: List[str],
+    score_per_tag: float = 0.25,
+    require_balanced: bool = True,
+    **kwargs: Any,
+) -> EvaluateResult:
+    """
+    Reward function that checks for presence of specific tags in response.
+
+    For each tag found in required_tags, adds score_per_tag to the score.
+    Optionally requires tags to be balanced (equal opening and closing tags).
+
+    Args:
+        messages: List of conversation messages
+        required_tags: List of tag names to check for (without < > brackets)
+        score_per_tag: Score to award per correctly found tag (default: 0.25)
+        require_balanced: If True, requires equal opening and closing tags
+        **kwargs: Additional arguments
+
+    Returns:
+        EvaluateResult with score based on tags found
+    """
+    if not messages or len(messages) == 0:
+        return EvaluateResult(
+            score=0.0,
+            reason="No messages provided",
+            metrics={"tag_count": MetricResult(score=0.0, is_score_valid=False, reason="No messages provided")},
+        )
+
+    response = messages[-1]
+
+    if response.role != "assistant" or not response.content:
+        return EvaluateResult(
+            score=0.0,
+            reason="No assistant response found or response has no content",
+            metrics={
+                "tag_count": MetricResult(
+                    score=0.0,
+                    is_score_valid=False,
+                    reason="Message not from assistant or has no content",
+                )
+            },
+        )
+    text: str = response.content
+
+    tag_metrics = {}
+    found_tags: Set[str] = set()
+    mismatched_tags: Set[str] = set()
+    total_found = 0
+
+    for tag in required_tags:
+        opening_pattern = f"<{tag}[^>]*>"
+        closing_pattern = f"</{tag}>"
+
+        opening_count = len(re.findall(opening_pattern, text))
+        closing_count = len(re.findall(closing_pattern, text))
+
+        if require_balanced:
+            is_found = opening_count > 0 and closing_count > 0 and opening_count == closing_count
+        else:
+            is_found = opening_count > 0 or closing_count > 0
+
+        is_balanced = opening_count == closing_count
+
+        if is_found:
+            found_tags.add(tag)
+            total_found += 1
+
+        if require_balanced and not is_balanced and (opening_count > 0 or closing_count > 0):
+            mismatched_tags.add(tag)
+
+        if require_balanced:
+            tag_score = 1.0 if (opening_count > 0 and closing_count > 0 and is_balanced) else 0.0
+            tag_success = opening_count > 0 and closing_count > 0 and is_balanced
+        else:
+            has_tags = opening_count > 0 or closing_count > 0
+            tag_score = 1.0 if has_tags else 0.0
+            tag_success = opening_count > 0 or closing_count > 0
+
+        tag_metrics[f"tag_{tag}"] = MetricResult(
+            score=tag_score,
+            is_score_valid=tag_success,
+            reason=_get_tag_reason(tag, opening_count, closing_count, require_balanced),
+        )
+
+    total_score = min(1.0, len(found_tags) * score_per_tag)
+
+    if require_balanced and mismatched_tags:
+        penalty = len(mismatched_tags) * score_per_tag
+        total_score = max(0.0, total_score - penalty)
+
+    success = len(found_tags) == len(required_tags) and (not require_balanced or not mismatched_tags)
+
+    reason = _get_overall_reason(required_tags, found_tags, mismatched_tags, require_balanced)
+    tag_metrics["overall"] = MetricResult(score=total_score, is_score_valid=success, reason=reason)
+
+    return EvaluateResult(score=total_score, reason=reason, metrics=tag_metrics, is_score_valid=success)
+
+
+def _get_tag_reason(tag: str, opening_count: int, closing_count: int, require_balanced: bool) -> str:
+    """Generate a descriptive reason for a tag's evaluation."""
+    if opening_count == 0 and closing_count == 0:
+        return f"Tag '{tag}' not found in response"
+    elif opening_count > 0 and closing_count == 0:
+        return f"Found {opening_count} opening <{tag}> tag(s) but no closing"
+    elif opening_count == 0 and closing_count > 0:
+        return f"Found {closing_count} closing </{tag}> tag(s) but no opening"
+    elif opening_count == closing_count:
+        return f"Found {opening_count} balanced '{tag}' tag(s)"
+    else:
+        if require_balanced:
+            return f"Unbalanced tags: {opening_count} opening vs " f"{closing_count} closing '{tag}' tags"
+        else:
+            return f"Found '{tag}' tags (unbalanced: {opening_count} opening, " f"{closing_count} closing)"
+
+
+def _get_overall_reason(
+    required_tags: List[str],
+    found_tags: Set[str],
+    mismatched_tags: Set[str],
+    require_balanced: bool,
+) -> str:
+    """Generate an overall reason for the evaluation."""
+    if not found_tags:
+        return "No required tags found in response"
+
+    missing_tags = set(required_tags) - found_tags
+
+    if not missing_tags and not mismatched_tags:
+        return f"All {len(required_tags)} required tags found and balanced"
+
+    reason_parts = []
+
+    if found_tags:
+        reason_parts.append(f"Found {len(found_tags)}/{len(required_tags)} required tags")
+
+    if missing_tags:
+        tags_str = ", ".join([f"'{tag}'" for tag in missing_tags])
+        reason_parts.append(f"Missing tags: {tags_str}")
+
+    if require_balanced and mismatched_tags:
+        tags_str = ", ".join([f"'{tag}'" for tag in mismatched_tags])
+        reason_parts.append(f"Unbalanced tags: {tags_str}")
+
+    return ". ".join(reason_parts)

eval_protocol/rl_processing.py

@@ -0,0 +1,82 @@
+from typing import Any, Dict, List, Optional, Union
+
+from eval_protocol.agent.models import StepData  # Internal StepData model
+
+# Assuming models are structured as planned
+from eval_protocol.models import EvaluateResult  # Extended EvaluateResult
+
+# Placeholder for actual Message type if needed for type hinting complex observation_data
+# from eval_protocol.models import Message
+
+
+class RLDataAligner:
+    """
+    Component responsible for aligning outputs from user reward functions
+    (EvaluateResult containing scores and/or base_rewards per step)
+    with the system's internal StepData representation. This prepares
+    the data for subsequent GiGPO (or other RL algorithm) advantage calculations.
+    """
+
+    def align_data_for_rl_processing(
+        self,
+        current_eval_result: EvaluateResult,
+        current_step_data_list: List[StepData],
+        rollout_id: str,  # For logging or if needed
+    ) -> List[StepData]:
+        """
+        Aligns the EvaluateResult (from user's reward function) with the
+        internal list of StepData for a single rollout.
+
+        Populates `StepData.base_reward` from `EvaluateResult.step_outputs.base_reward`.
+        Associates `EvaluateResult.score` with the rollout for GiGPO A_E calculation.
+        (Association of final_score might happen by returning it alongside, or
+         by the caller managing it). For now, this function focuses on base_reward.
+
+        Args:
+            current_eval_result: The EvaluateResult from the user's reward function for this rollout.
+            current_step_data_list: The list of StepData objects collected by RLRolloutWorker.
+            rollout_id: Identifier for the current rollout.
+
+        Returns:
+            The list of StepData objects, with `base_reward` populated.
+            The `final_score` from current_eval_result should be handled by the caller
+            for GiGPO A_E calculation.
+        """
+
+        # Store final_score (for GiGPO A_E) - The caller will handle this.
+        # This function's primary job is to populate base_rewards in StepData.
+
+        if current_eval_result.step_outputs:
+            # Create a dictionary for quick lookup of user-defined step rewards
+            user_step_rewards_map: Dict[Union[int, str], float] = {
+                step_out.step_index: step_out.base_reward for step_out in current_eval_result.step_outputs
+            }
+
+            for s_data in current_step_data_list:
+                # --- Critical Mapping Logic ---
+                # Strategy: Use 'assistant_turn_index' stored in StepData.step_info
+                # by RLRolloutWorker. User's StepOutput.step_index should match this.
+                # This assumes RLRolloutWorker adds this info.
+                user_defined_step_idx = s_data.step_info.get("assistant_turn_index")
+
+                if user_defined_step_idx is not None:
+                    if user_defined_step_idx in user_step_rewards_map:
+                        s_data.base_reward = user_step_rewards_map[user_defined_step_idx]
+                    else:
+                        # No base reward provided by user for this specific system step.
+                        # s_data.base_reward remains None (or could be a default).
+                        pass
+                else:
+                    # RLRolloutWorker did not provide 'assistant_turn_index' for this StepData,
+                    # or the mapping key in step_info is different.
+                    # This indicates a potential issue in RLRolloutWorker or mapping strategy.
+                    pass
+        else:
+            # No step_outputs provided by the user. Base rewards will remain None.
+            pass
+
+        return current_step_data_list
+
+    # TODO (Future): Consider a batch version if performance becomes an issue,
+    # but the core logic per rollout remains the same.
+    # def align_batch_data_for_rl_processing(...)