# 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.
import json
import re
from abc import ABC, abstractmethod
from typing import Any, Optional, Union
from tqdm.auto import tqdm
from typing_extensions import Self
from oumi.core.configs import InferenceConfig, JudgeConfig
from oumi.core.inference import BaseInferenceEngine
from oumi.core.types.conversation import Conversation, Message, Role, TemplatedMessage
from oumi.utils.logging import logger
[docs]
class BaseJudgeOutput(ABC, TemplatedMessage):
raw_judgement: Optional[str] = None
[docs]
@classmethod
def from_xml_output(cls, raw_judgement: Optional[str]) -> Optional[Self]:
"""Parses the judgement from XML-like tags in the raw output.
Args:
raw_judgement: The raw judgement string to parse.
Returns:
Optional[Self]: An instance of the class with parsed attributes,
or None if parsing fails.
"""
if not raw_judgement:
return None
attributes = {}
# Regex pattern to match XML-like tags and their content
# Captures the tag name in group 1 and the content between tags in group 2
# For example, "<label>True</label>" would match as ("label", "True")
pattern = r"<(\w+)>(.*?)</\1>"
matches = re.findall(pattern, raw_judgement, re.DOTALL)
for attr_name, attr_value in matches:
attributes[attr_name] = attr_value.strip()
return cls(**attributes, raw_judgement=raw_judgement)
[docs]
@classmethod
def from_json_output(cls, raw_judgement: Optional[str]) -> Optional[Self]:
"""Parses the judgement from JSON."""
if not raw_judgement:
return None
try:
judgement_data = json.loads(raw_judgement)
return cls(**judgement_data, raw_judgement=raw_judgement)
except json.JSONDecodeError:
return None
@property
def label(self):
"""Convert the judgement to a boolean or Likert scale label.
This method should be overridden by subclasses to provide the actual
conversion logic.
"""
return self.raw_judgement
@property
def fields(self):
"""Return the fields of the judgement."""
fields = self.model_dump()
fields.pop("raw_judgement", None)
fields.pop("template", None)
fields.pop("role", None)
return fields
[docs]
class BaseJudge(ABC):
def __init__(
self,
config: JudgeConfig,
inference_engine: Optional[BaseInferenceEngine] = None,
):
"""Initialize the Judge."""
self._config = config
self._attributes = config.attributes
if len(config.attributes) < 1:
raise ValueError(
"At least one attribute must be specified in the judge configuration."
)
if inference_engine is not None:
logger.debug("Using provided inference engine.")
self.inference_engine = inference_engine
else:
logger.debug("Initializing inference engine.")
self.inference_engine = self._create_inference_engine(config)
[docs]
def judge(
self,
raw_inputs: Union[list[Conversation], list[dict], list[Message]],
) -> list[dict[str, BaseJudgeOutput]]:
"""Judge the given conversations."""
# Convert the raw user inputs into a list of JudgeInput classes
# A JudgeInput is the unit of what needs to be judged, and could be a
# prompt, request/answer pair or a full conversation
judge_inputs = []
for raw_input in raw_inputs:
if isinstance(raw_input, dict):
judge_input = self._transform_dict_input(raw_input)
elif isinstance(raw_input, TemplatedMessage):
judge_input = raw_input
elif isinstance(raw_input, Conversation):
judge_input = self._transform_conversation_input(raw_input)
else:
raise ValueError(f"Unsupported conversation type: {type(raw_input)}")
judge_inputs.append(judge_input)
results = {}
for attribute_name in self._attributes.keys():
# Generate the full judging prompt for each attribute
# This includes the judge system prompt, and any few shot examples
# That are included in the judge config.
judgement_prompts = [
self.build_judgement_prompt(judge_input, attribute_name=attribute_name)
for judge_input in tqdm(judge_inputs)
]
# Run inference for the attribute's prompt
# We batch the inference for a single attribute together to maximally
# benefit from kv prefix caching (system prompt, few shot examples)
raw_judgements = self._infer(judgement_prompts)
# Parse the raw judge output (a string) into a JudgeOutput object
judgements = []
for conversation in raw_judgements:
judgement = conversation.messages[-1].content
parsed_judgement = self._transform_model_output(judgement)
judgements.append(
{
"raw_judgement": judgement,
"fields": parsed_judgement.fields,
"label": parsed_judgement.label,
}
)
results[attribute_name] = judgements
# Results are now in the form
# {attribute: judgements for attribute in attributes}
# Transform to
# [{attribute: judgement} for judgement in judgements]
outputs = []
for idx in range(len(raw_inputs)):
output_dict = {}
for attribute_name in self._attributes.keys():
output_dict[attribute_name] = results[attribute_name][idx]
outputs.append(output_dict)
return outputs
[docs]
def build_judgement_prompt(
self, judge_input: Message, attribute_name: Optional[str]
) -> Conversation:
"""Generate judge prompts for a dataset."""
if attribute_name is None:
if len(self._attributes) > 0:
raise ValueError(
"attribute_name must be specified when there are multiple"
" attributes to judge."
)
else:
# If there's only one attribute, use it
attribute_name = next(iter(self._attributes))
if attribute_name not in self._attributes:
raise KeyError(
f"Attribute '{attribute_name}' not found in config.attributes"
)
attribute = self._attributes[attribute_name]
messages = attribute.messages.copy()
messages.append(Message(content=judge_input.content, role=Role.USER))
return Conversation(
messages=messages,
metadata={
"judge_attribute_name": attribute_name,
},
)
def _infer(self, conversations: list[Conversation]) -> list[Conversation]:
"""Judge a single attribute."""
metadatas = [convo.metadata for convo in conversations]
# Wrap the generation params in an inference config for inference.
# Only the generations params are used by the inference engine.
inference_config = InferenceConfig(
model=self._config.model,
generation=self._config.generation,
remote_params=self._config.remote_params,
)
responses = self.inference_engine.infer(
input=conversations, inference_config=inference_config
)
assert len(responses) == len(metadatas)
for response, metadata in zip(responses, metadatas):
response.metadata.update(metadata)
return responses
def _create_inference_engine(self, config: JudgeConfig) -> BaseInferenceEngine:
"""Create the inference engine."""
from oumi.builders.inference_engines import build_inference_engine
return build_inference_engine(
engine_type=config.engine,
model_params=config.model,
remote_params=config.remote_params,
)
@abstractmethod
def _transform_conversation_input(self, conversation: Conversation) -> Message:
raise NotImplementedError
@abstractmethod
def _transform_dict_input(self, raw_input: dict[str, Any]) -> Message:
raise NotImplementedError
@abstractmethod
def _transform_model_output(self, model_output) -> BaseJudgeOutput:
raise NotImplementedError
