Repository: oumi-ai/oumi
File: src/oumi/core/evaluation/evaluator.py (Python)

Code before:
            evaluation_fn = self._get_custom_evaluation_fn(evaluation_fn_name)
            custom_kwargs = self._merge_kwargs(kwargs, task_params.eval_kwargs)
            self._validate_custom_kwargs(
                custom_kwargs=custom_kwargs,
                evaluation_fn=evaluation_fn,
                evaluation_fn_name=evaluation_fn_name,
            )
            self._add_reserved_keys_into_custom_kwargs(
                custom_kwargs=custom_kwargs,
                evaluation_fn=evaluation_fn,
                task_params=task_params,
                config=config,
            )
            evaluation_output = evaluation_fn(**custom_kwargs)

            if isinstance(evaluation_output, EvaluationResult):
                evaluation_result = evaluation_output
            elif isinstance(evaluation_output, dict):
                evaluation_result = EvaluationResult(
                    task_name=task_params.task_name,
                    task_result={"results": {task_params.task_name: evaluation_output}},
                )
            else:
                raise ValueError(
                    f"The custom evaluation function `{task_params.task_name}` must "
                    "return a `dict` object but, instead, it's currently returning "
                    f"an object of type `{type(evaluation_output)}`. Please ensure "
                    "that the function returns the correct object."
                )
        else:
            raise ValueError(f"Unknown evaluation backend: {evaluation_backend}")

        # Calculate the elapsed time for the evaluation run.
        evaluation_result.elapsed_time_sec = int(time.time() - start_time)
        evaluation_result.start_time = start_time_str

        # Save the output, if an output directory has been provided.
        if config.output_dir and is_world_process_zero():
            self.save_output(
                task_params=task_params,
                evaluation_result=evaluation_result,
                base_output_dir=config.output_dir,
                config=config,
            )
        return evaluation_result

    def save_output(
        self,
        task_params: EvaluationTaskParams,
        evaluation_result: EvaluationResult,
        base_output_dir: str,

Review comment (suggestion):
I know you don't want to advertise the EvaluationResult option, but without mentioning it, this error isn't correct. Could you still mention it here? You could highlight that the dict is preferred though.

Code after:
            evaluation_fn = self._get_custom_evaluation_fn(evaluation_fn_name)
            custom_kwargs = self._merge_kwargs(kwargs, task_params.eval_kwargs)
            self._validate_custom_kwargs(
                custom_kwargs=custom_kwargs,
                evaluation_fn=evaluation_fn,
                evaluation_fn_name=evaluation_fn_name,
            )
            self._add_reserved_keys_into_custom_kwargs(
                custom_kwargs=custom_kwargs,
                evaluation_fn=evaluation_fn,
                task_params=task_params,
                config=config,
            )
            evaluation_output = evaluation_fn(**custom_kwargs)

            if isinstance(evaluation_output, EvaluationResult):
                evaluation_result = evaluation_output
            elif isinstance(evaluation_output, dict):
                evaluation_result = EvaluationResult(
                    task_name=task_params.task_name,
                    task_result={"results": {task_params.task_name: evaluation_output}},
                )
            else:
                raise ValueError(
                    f"The custom evaluation function `{task_params.task_name}` must "
                    "return either a `dict` or an `EvaluationResult` object, but it is "
                    f"currently returning an object of type `{type(evaluation_output)}`"
                    ". Please ensure that the function returns the correct object."
                )
        else:
            raise ValueError(f"Unknown evaluation backend: {evaluation_backend}")

        # Calculate the elapsed time for the evaluation run.
        evaluation_result.elapsed_time_sec = int(time.time() - start_time)
        evaluation_result.start_time = start_time_str

        # Save the output, if an output directory has been provided.
        if config.output_dir and is_world_process_zero():
            self.save_output(
                task_params=task_params,
                evaluation_result=evaluation_result,
                base_output_dir=config.output_dir,
                config=config,
            )
        return evaluation_result

    def save_output(
        self,
        task_params: EvaluationTaskParams,
        evaluation_result: EvaluationResult,
        base_output_dir: str,