# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing_extensions import override
from oumi.core.configs.judge_config import JudgeConfig
from oumi.core.configs.params.judge_params import (
JudgeOutputType,
JudgeParams,
JudgeResponseFormat,
)
from oumi.judges.base_judge import (
BaseJudge,
JudgeOutputField,
)
from oumi.judges.judge_utils import (
build_judgment_field_schema,
describe_judgment_options,
)
# Expected field/key names in the judge's output.
EXPLANATION_KEY = "explanation"
JUDGMENT_KEY = "judgment"
# Prompt suffix: describing to the judge how to format its response (XML, JSON, or RAW).
XML_SUFFIX = (
"\n\nProvide your response in XML format only. Include your judgment enclosed "
"within <{judgment_key}> and </{judgment_key}> tags. {judgment_options}Do not "
"include any text outside the XML. Ensure that all tags are properly closed and "
"that the XML is well-formed."
)
XML_SUFFIX_WITH_EXPLANATION = (
"\n\nProvide your response in XML format only. Begin with an explanation "
"justifying your judgment, enclosed within <{explanation_key}> and "
"</{explanation_key}> tags. Follow this with your judgment, enclosed within "
"<{judgment_key}> and </{judgment_key}> tags. {judgment_options}Do not include any "
"text outside the XML. Ensure that all tags are properly closed and that the XML "
"is well-formed."
)
JSON_SUFFIX = (
"\n\nProvide your response in JSON format only. Include your judgment as the value "
"of a single key named '{judgment_key}'. {judgment_options}Do not include any "
"text outside the JSON. Ensure the JSON is properly formatted and valid."
)
JSON_SUFFIX_WITH_EXPLANATION = (
"\n\nProvide your response in JSON format only. Begin with an explanation "
"justifying your judgment, using the key '{explanation_key}'. Then include your "
"judgment using the key '{judgment_key}'. {judgment_options}Do not include any "
"text outside the JSON. Ensure the JSON is properly formatted and valid."
)
RAW_SUFFIX_WITH_EXPLANATION = (
"\n\nExplain your reasoning before providing your judgment."
)
[docs]
class SimpleJudge(BaseJudge):
"""Judge class for evaluating outputs based on a given configuration."""
def __init__(
self,
judge_config: JudgeConfig | str,
):
"""Initialize the Judge.
Args:
judge_config: JudgeConfig object or a path to a judge configuration file.
Contains both judge parameters and inference configuration.
"""
if isinstance(judge_config, str):
judge_config = JudgeConfig.from_path(judge_config)
self._judge_params = judge_config.judge_params
self._judge_params.replace_template_variables()
self._inference_config = judge_config.inference_config
# Create output fields based on judge configuration
output_fields = []
if self._judge_params.include_explanation:
output_fields.append(self._create_explanation_output_field())
output_fields.append(self._create_judgment_output_field(self._judge_params))
# Generate an inference engine from inference config
if self._inference_config is None:
raise ValueError(
"inference_config must be provided in JudgeConfig for SimpleJudge. "
"Please ensure your JudgeConfig includes a valid inference_config."
)
use_schema = (
self._judge_params.response_format == JudgeResponseFormat.JSON
and self._judge_params.use_guided_decoding
)
inference_engine = self._create_inference_engine(
inference_config=self._inference_config,
response_schema=self._build_response_schema() if use_schema else None,
)
# Append format suffix to system instruction if it exists
system_instruction = self._judge_params.system_instruction
if system_instruction:
system_instruction = f"{system_instruction}{self._get_format_suffix()}"
# Get set of prompt template placeholders
prompt_template_placeholders_set = (
set(self._judge_params.prompt_template_placeholders)
if self._judge_params.prompt_template_placeholders
else self._judge_params.get_placeholders()
)
super().__init__(
prompt_template=self._judge_params.prompt_template,
prompt_template_placeholders=prompt_template_placeholders_set,
system_instruction=system_instruction,
example_field_values=self._judge_params.examples,
response_format=self._judge_params.response_format,
output_fields=output_fields,
inference_engine=inference_engine,
)
@override
def _build_judgment_prompt(self, judge_input: dict[str, str]) -> str:
"""Generate judge prompts using the template."""
prompt_content = super()._build_judgment_prompt(judge_input)
# Only append format suffix to judgment prompt if no system instruction exists
# (otherwise it was already appended to system instruction in __init__)
if not self._judge_params.system_instruction:
prompt_content += self._get_format_suffix()
return prompt_content
def _get_format_suffix(self) -> str:
"""Get the appropriate format suffix based on response format and explanation.
Returns:
Format-specific instruction suffix to append to prompts
"""
response_format = self._judge_params.response_format
include_explanation = self._judge_params.include_explanation
# Describe the expected judgment options to the judge
judgment_options = describe_judgment_options(
judgment_type=self._judge_params.judgment_type,
judgment_scores=self._judge_params.judgment_scores,
)
# Describe the expected response format to the judge
if response_format == JudgeResponseFormat.XML:
suffix = XML_SUFFIX_WITH_EXPLANATION if include_explanation else XML_SUFFIX
elif response_format == JudgeResponseFormat.JSON:
suffix = (
JSON_SUFFIX_WITH_EXPLANATION if include_explanation else JSON_SUFFIX
)
elif response_format == JudgeResponseFormat.RAW:
suffix = RAW_SUFFIX_WITH_EXPLANATION if include_explanation else ""
else:
suffix = ""
return suffix.format(
judgment_key=JUDGMENT_KEY,
explanation_key=EXPLANATION_KEY,
judgment_options=judgment_options,
)
def _create_judgment_output_field(self, params: JudgeParams) -> JudgeOutputField:
"""Create the main judgment output field."""
# dict is invariant, so dict[str, float] is not assignable to
# dict[str, float | None]. A simple judge never scores a label None, so
# copying into the wider dict is sound.
field_scores: dict[str, float | None] | None = (
dict(params.judgment_scores) if params.judgment_scores else None
)
return JudgeOutputField(
field_key=JUDGMENT_KEY,
field_type=params.judgment_type,
field_scores=field_scores,
)
def _create_explanation_output_field(self) -> JudgeOutputField:
"""Create the explanation output field."""
return JudgeOutputField(
field_key=EXPLANATION_KEY,
field_type=JudgeOutputType.TEXT,
field_scores=None,
)
def _build_response_schema(self) -> dict:
"""JSON schema describing the expected judge response."""
properties: dict[str, dict] = {}
# Add explanation field, if required (Note: MUST BE the first field to add)
if self._judge_params.include_explanation:
properties[EXPLANATION_KEY] = {"type": "string"}
# Add judgment field
properties[JUDGMENT_KEY] = build_judgment_field_schema(
judgment_type=self._judge_params.judgment_type,
judgment_scores=self._judge_params.judgment_scores,
)
return {
"type": "object",
"properties": properties,
"required": list(properties.keys()),
"additionalProperties": False,
}
