afnio.cognitive.modules.lm_judge_evaluator

class LMJudgeEvaluator(Module):
    """
    Evaluates predictions using a language model (LM) as the judge.

    This module leverages the [`LMJudgeEvaluator`][afnio.autodiff.evaluator.LMJudgeEvaluator]
    operation from `afnio.autodiff.evaluator` to perform model-based evaluations.
    The `forward` method accepts a list of `messages` that construct the evaluation
    prompt, with optional `inputs` to dynamically fill placeholders within message
    templates. A `prediction` is compared against a `target` (optional) to generate
    a `score` and an `explanation`.

    When processing a batch of predictions and targets, `reduction_fn` function
    aggregates individual scores (e.g., using [`sum`][sum] to compute a total score). The
    `reduction_fn_purpose` parameter is a brief description of the aggregation's purpose
    (e.g., `"summation"`). If aggregation is not desired, set `reduction_fn` and
    `reduction_fn_purpose` to `None`. The `success_fn` checks if all evaluations are
    successful, allowing the `backward` pass to skip unnecessary gradient computations.

    This module supports both evaluation (`eval_mode=True`) and optimization
    (`eval_mode=False`) modes.

    The `forward_model_client` specifies the LM responsible for evaluation, while
    `completion_args` allows customization of generation parameters like temperature,
    max tokens, and seed.

    Examples:
        >>> from afnio import cognitive as cog
        >>> from afnio.models.openai import OpenAI
        >>> from afnio import set_backward_model_client
        >>> fwd_model_client = OpenAI()
        >>> fwd_model_args = {"model": "gpt-4o", "temperature": 0.5}
        >>> set_backward_model_client("openai/gpt-4o")
        >>> class Evaluator(cog.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.judge = cog.LMJudgeEvaluator()
        ...     def forward(self, fwd_model, messages, prediction, target, inputs, **completion_args):
        ...         return self.judge(fwd_model, messages, prediction, target, inputs, **completion_args)
        >>> task = afnio.Variable(
        ...     "Evaluate if the translation is {metric}.",
        ...     role="evaluation task",
        ...     requires_grad=True
        ... )
        >>> format = afnio.Variable(
        ...     "Provide 'score' (true/false) and 'explanation' in JSON.",
        ...     role="output format"
        ... )
        >>> metric = afnio.Variable(["accurate", "accurate"], role="metric")
        >>> user = afnio.Variable(
        ...     "<PREDICTION>{prediction}</PREDICTION><TARGET>{target}</TARGET>",
        ..      role="user query"
        ... )
        >>> prediction = afnio.Variable(
        ...     ["Hola Mundo", "Salve a tutti"],
        ...     role="translated text",
        ...     requires_grad=True
        ... )
        >>> target = ["Ciao Mondo", "Salve a tutti"]
        >>> messages = [
        ...     {"role": "system", "content": [task, format]},
        ...     {"role": "user", "content": [user]},
        ... ]
        >>> eval = Evaluator()
        >>> score, explanation = eval(
        ...     fwd_model_client,
        ...     messages,
        ...     prediction,
        ...     target,
        ...     inputs={"metric": metric},
        ...     reduction_fn=sum,
        ...     reduction_fn_purpose="summation",
        ...     **fwd_model_args
        ... )
        >>> print(score.data)
        1
        >>> print(explanation.data)
        'The evaluation function, designed using an LM as the judge, compared the <DATA> fields of the predicted variable and the target variable across all samples in the batch. These scores were then aggregated using the reduction function 'summation', resulting in a final aggregated score: 1.'
        >>> explanation.backward()
        >>> system.grad[0].data
        'The translated text should be in Italian.'

    Raises:
        RuntimeError: If the LM response to generate the evaluation `score` and
            `explanation` cannot be parsed as valid JSON.
        TypeError: If the types of `forward_model_client`, `messages`, `prediction`,
            `target`, `inputs`, `success_fn`, `reduction_fn`, `reduction_fn_purpose`,
            or `eval_mode` are not as expected.
        ValueError: If the lengths of `prediction.data` and `target` (or `target.data`,
            when `target` is a `Variable`) do not match when both are lists, or if
            `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is an empty string,
            or if `inputs` contains keys that conflict with `prediction` or `target`.

    See Also:
        [`afnio.autodiff.evaluator.LMJudgeEvaluator`][afnio.autodiff.evaluator.LMJudgeEvaluator]
        for the underlying operation.
    """  # noqa: E501

    forward_model_client: Optional[ChatCompletionModel]
    messages: MultiTurnMessages
    success_fn: Optional[Callable[[List[Any]], bool]]
    reduction_fn: Optional[Callable[[List[Any]], Any]]
    reduction_fn_purpose: Optional[Union[str, Variable]]
    eval_mode: Union[bool, Variable]
    completion_args: Dict[str, Any]

    def __init__(self):
        super().__init__()

        self.register_model("forward_model_client", None)
        self.register_chat("messages", None)
        self.register_function("success_fn", None)
        self.register_function("reduction_fn", None)
        self.register_buffer("reduction_fn_purpose", None)
        self.register_buffer("eval_mode", None)
        self.register_completion_config("completion_args", None)

    def forward(
        self,
        forward_model_client: Optional[ChatCompletionModel],
        messages: MultiTurnMessages,
        prediction: Variable,
        target: Optional[Union[str, List[str], Variable]] = None,
        inputs: Optional[Dict[str, Union[str, Variable]]] = None,
        success_fn: Optional[Callable[[List[Any]], bool]] = None,
        reduction_fn: Optional[Callable[[List[Any]], Any]] = sum,
        reduction_fn_purpose: Optional[Union[str, Variable]] = "summation",
        eval_mode: Union[bool, Variable] = True,
        **completion_args,
    ) -> Tuple[Variable, Variable]:
        """
        Forward pass for the LM Judge evaluator function.

        Warning:
            Users should not call this method directly. Instead, they should call the
            module instance itself, which will internally invoke this `forward` method.

        Args:
            forward_model_client: The LM model client used for the forward
                pass evaluation.
            messages: A list of messages that compose the prompt/context for the LM.
                Each message is a dictionary with a `"role"` (e.g., `"system"`,
                `"user"`, `"assistant"`) and a `"content"` field, which is a list of
                `Variable` objects. The `Variable` objects in the `"content"` can
                contain placeholders (e.g., `{prediction}`, `{target}`) that will be
                populated with the corresponding values from the `inputs` dictionary.
            prediction: The predicted variable to evaluate, which can have scalar or
                list [`data`][afnio.Variable.data] (supporting both individual and
                batch processing).
            target: The target (ground truth) to compare against, which can be a string,
                a list of strings, or a `Variable`. Optional if the evaluation does not
                require a target and only relies on the correctness of the LM Judge's
                assessment of the `prediction`.
            inputs: A dictionary mapping placeholder names to their corresponding
                values, which can be strings or `Variable` instances. These values
                will be used to populate the placeholders in the `messages` content
                before sending the prompt to the LM. For example, if a message
                `"content"` field contains the placeholder `{color}`, the `inputs`
                dictionary should have a key `"color"` with the value to substitute
                in the prompt. Optional if there are no placeholders in the messages or
                if all placeholders are directly related to `prediction` and `target`.
            success_fn: A user-defined function that takes the list of scores returned
                by the LM Judge and returns `True` if all predictions are considered
                successful, or `False` otherwise.
            reduction_fn: An optional function to aggregate scores across a batch of
                predictions and targets. If `None`, no aggregation is applied.
            reduction_fn_purpose: A brief description of the purpose of `reduction_fn`,
                used by the autodiff engine to generate explanations. Required if
                `reduction_fn` is provided.
            eval_mode: Indicates the evaluation mode. If `True`, the `backward` pass
                will compute gradients for the `prediction` variable only. If `False`,
                the `backward` pass will compute gradients for the `messages` and
                `inputs`, allowing optimization of the evaluator itself or alignment
                with human evaluation datasets.
            **completion_args: Additional keyword arguments to pass to the LM model
                client's `chat` method, such as temperature, max tokens, or seed values,
                to customize the LLM's behavior during the evaluation.

        Returns:
            score: A variable containing the evaluation score(s),
                or their aggregation if `reduction_fn` is provided.
            explanation: A variable containing the explanation(s) of the evaluation,
                or their aggregation if `reduction_fn` is provided.

        Raises:
            RuntimeError: If the LM response to generate the evaluation `score` and
                `explanation` cannot be parsed as valid JSON.
            TypeError: If the types of `forward_model_client`, `messages`, `prediction`,
                `target`, `inputs`, `success_fn`, `reduction_fn`,
                `reduction_fn_purpose`, or `eval_mode` are not as expected.
            ValueError: If the lengths of `prediction.data` and `target` (or
                `target.data`, when `target` is a `Variable`) do not match when both are
                lists, or if `reduction_fn_purpose` (or `reduction_fn_purpose.data`) is
                an empty string, or if `inputs` contains keys that conflict with
                `prediction` or `target`.
        """
        self.forward_model_client = forward_model_client
        self.messages = messages
        self.success_fn = success_fn
        self.reduction_fn = reduction_fn
        self.reduction_fn_purpose = (
            None
            if reduction_fn_purpose is None
            else (
                reduction_fn_purpose
                if isinstance(reduction_fn_purpose, Variable)
                else Variable(reduction_fn_purpose)
            )
        )
        self.eval_mode = (
            eval_mode if isinstance(eval_mode, Variable) else Variable(eval_mode)
        )
        self.completion_args = completion_args
        return LMJudgeEvaluatorOp.apply(
            self.forward_model_client,
            self.messages,
            prediction,
            target,
            inputs,
            self.success_fn,
            self.reduction_fn,
            self.reduction_fn_purpose,
            self.eval_mode,
            **self.completion_args,
        )