# 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 import Union
from oumi.core.configs import (
GenerationParams,
JudgeConfig,
ModelParams,
RemoteParams,
)
from oumi.core.configs.inference_config import InferenceEngineType
from oumi.core.configs.judge_config import JudgeAttribute
from oumi.core.registry import register_judge
from oumi.judges.oumi_judge import OumiJudgeInput, OumiJudgeOutput
from oumi.utils.io_utils import get_oumi_root_directory
[docs]
@register_judge("oumi/v1_xml_claude_sonnet")
def oumi_v1_xml_claude_sonnet_judge() -> JudgeConfig:
"""Returns a JudgeConfig for the Oumi v1 XML Anthropic judge.
This function creates and returns a JudgeConfig object for the Oumi V1 Judge, which
uses Claude Sonnet as a judge, with inputs and outputs in XML format.
Returns:
JudgeConfig: A configuration object for the Oumi v1 XML Anthropic judge.
Note:
This judge uses the Anthropic API, so the ANTHROPIC_API_KEY environment
variable must be set with a valid API key.
"""
judges_directory = get_oumi_root_directory() / "judges" / "oumi_v1"
attribute_names = ["helpful", "honest", "safe"]
attributes = {
attribute: JudgeAttribute[Union[OumiJudgeInput, OumiJudgeOutput]].load(
str(judges_directory / f"{attribute}.json")
)
for attribute in attribute_names
}
config = JudgeConfig(
attributes=attributes,
model=ModelParams(
model_name="claude-3-5-sonnet-20240620",
),
generation=GenerationParams(
max_new_tokens=1024,
),
remote_params=RemoteParams(
api_url="https://api.anthropic.com/v1/messages",
api_key_env_varname="ANTHROPIC_API_KEY",
max_retries=3,
),
engine=InferenceEngineType.ANTHROPIC,
)
return config
[docs]
@register_judge("oumi/v1_xml_local")
def oumi_v1_xml_local_judge() -> JudgeConfig:
"""Returns a JudgeConfig for the Oumi v1 XML local judge.
Returns:
JudgeConfig: A configuration object for the Oumi v1 XML local judge.
Note:
This judge uses a local GGUF model file for inference.
"""
judges_directory = get_oumi_root_directory() / "judges" / "oumi_v1"
attribute_names = ["helpful", "honest", "safe"]
attributes = {
attribute: JudgeAttribute[Union[OumiJudgeInput, OumiJudgeOutput]].load(
str(judges_directory / f"{attribute}.json")
)
for attribute in attribute_names
}
config = JudgeConfig(
attributes=attributes,
model=ModelParams(model_name="Qwen/Qwen2-0.5B-Instruct-GGUF"),
engine=InferenceEngineType.LLAMACPP,
generation=GenerationParams(max_new_tokens=1024, temperature=0.0),
)
return config
[docs]
@register_judge("oumi/v1_xml_gpt4o")
def oumi_v1_xml_gpt4o_judge() -> JudgeConfig:
"""Returns a JudgeConfig for the Oumi v1 XML GPT-4 judge.
This function creates and returns a JudgeConfig object for the Oumi V1 Judge, which
uses GPT-4 as a judge, with inputs and outputs in XML format.
Returns:
JudgeConfig: A configuration object for the Oumi v1 XML GPT-4 judge.
Note:
This judge uses the OpenAI API, so the OPENAI_API_KEY environment
variable must be set with a valid API key.
"""
judges_directory = get_oumi_root_directory() / "judges" / "oumi_v1"
attribute_names = ["helpful", "honest", "safe"]
attributes = {
attribute: JudgeAttribute[Union[OumiJudgeInput, OumiJudgeOutput]].load(
str(judges_directory / f"{attribute}.json")
)
for attribute in attribute_names
}
config = JudgeConfig(
attributes=attributes,
model=ModelParams(model_name="gpt-4o-2024-08-06"),
engine=InferenceEngineType.REMOTE,
generation=GenerationParams(
max_new_tokens=1024,
temperature=0.0,
),
remote_params=RemoteParams(
api_url="https://api.openai.com/v1/chat/completions",
api_key_env_varname="OPENAI_API_KEY",
max_retries=3,
),
)
return config
@register_judge("oumi/v1_xml_unit_test")
def unit_test_judge():
"""Tiny judge for unit testing.
Do not use this judge for anything serious as it returns random results.
"""
attribute_path = (
get_oumi_root_directory() / "judges" / "test_judge" / "helpful.json"
)
attribute = JudgeAttribute[Union[OumiJudgeInput, OumiJudgeOutput]].load(
str(attribute_path)
)
config = JudgeConfig(
attributes={"helpful": attribute},
engine=InferenceEngineType.NATIVE,
model=ModelParams(
model_name="gpt2",
tokenizer_pad_token="</s>",
chat_template="gpt2",
),
generation=GenerationParams(
max_new_tokens=128,
temperature=0.0,
),
)
return config
