# 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
from pathlib import Path
from typing import TYPE_CHECKING, Annotated, Optional
import jsonlines
import typer
from rich.table import Table
from oumi.cli import cli_utils
from oumi.utils.io_utils import load_jsonlines
if TYPE_CHECKING:
from oumi.core.configs import InferenceConfig, JudgeConfig
def _load_judge_config(config: str, extra_args: list[str]) -> "JudgeConfig":
from oumi.core.registry import REGISTRY
judge_config_builder = REGISTRY.get_judge_config(config)
if judge_config_builder:
if extra_args:
typer.echo(
"For consistent judge results, a named judge config cannot be "
f"overridden with extra arguments. Got: {extra_args}. "
"Please register a new named judge config, or provide a path to a "
"judge config file."
)
raise typer.Exit(code=1)
return judge_config_builder()
if not Path(config).exists():
typer.echo(f"Config file not found: '{config}'")
raise typer.Exit(code=1)
return JudgeConfig.from_yaml_and_arg_list(config, extra_args)
[docs]
def dataset(
ctx: typer.Context,
config: Annotated[
str, typer.Option(*cli_utils.CONFIG_FLAGS, help="Path to the judge config file")
],
dataset_name: Annotated[
Optional[str], typer.Option(help="Name of the dataset from the registry")
] = None,
dataset_subset: Annotated[
Optional[str], typer.Option(help="Subset of the dataset to use, if applicable")
] = None,
dataset_split: Annotated[
Optional[str], typer.Option(help="Split of the dataset to use.")
] = "train",
output_file: Annotated[
Optional[str], typer.Option(help="Path to the output file (jsonl)")
] = None,
level: cli_utils.LOG_LEVEL_TYPE = None,
):
"""Judge a dataset."""
# Delayed imports
from oumi import judge_dataset
from oumi.core.registry import REGISTRY
# End imports
if not dataset_name:
typer.echo("Dataset name is required.")
raise typer.Exit(code=1)
# Load the judge config
extra_args = cli_utils.parse_extra_cli_args(ctx)
config = str(
cli_utils.resolve_and_fetch_config(
config,
)
)
judge_config = _load_judge_config(config, extra_args)
# Load the dataset class from the registry
dataset_class = REGISTRY.get_dataset(dataset_name, subset=dataset_subset)
if dataset_class is None:
typer.echo(f"Dataset '{dataset_name}' not found in registry.")
raise typer.Exit(code=1)
dataset = dataset_class(
split=dataset_split,
subset=dataset_subset,
)
# Judge the dataset
results = judge_dataset(judge_config, dataset=dataset)
# Save the results
if output_file:
typer.echo(f"Saving results to {output_file}")
with jsonlines.open(output_file, mode="w") as writer:
writer.write_all(results)
else:
table = Table(
title="Judge Results",
title_style="bold magenta",
show_edge=False,
show_lines=True,
)
table.add_column("Judgements", style="green")
for result in results:
table.add_row(json.dumps(result))
cli_utils.CONSOLE.print(table)
[docs]
def conversations(
ctx: typer.Context,
config: Annotated[
str, typer.Option(*cli_utils.CONFIG_FLAGS, help="Path to the judge config file")
],
input_file: Annotated[
Optional[str], typer.Option(help="Path to the input file (jsonl)")
] = None,
output_file: Annotated[
Optional[str], typer.Option(help="Path to the output file (jsonl)")
] = None,
level: cli_utils.LOG_LEVEL_TYPE = None,
):
"""Judge a list of conversations."""
extra_args = cli_utils.parse_extra_cli_args(ctx)
config = str(
cli_utils.resolve_and_fetch_config(
config,
)
)
# Delayed imports
from oumi import judge_conversations
from oumi.core.types.conversation import Conversation
# End imports
# Load the judge config
judge_config = _load_judge_config(config, extra_args)
# Load the conversations from the input file
if not input_file:
typer.echo("Input file is required.")
raise typer.Exit(code=1)
input_data = load_jsonlines(input_file)
conversations = [Conversation.from_dict(conv) for conv in input_data]
# Judge the conversations
results = judge_conversations(judge_config, judge_inputs=conversations)
# Save the results
if output_file:
typer.echo(f"Saving results to {output_file}")
with jsonlines.open(output_file, mode="w") as writer:
writer.write_all(results)
else:
table = Table(
title="Judge Results",
title_style="bold magenta",
show_edge=False,
show_lines=True,
)
table.add_column("Judgements", style="green")
for result in results:
table.add_row(json.dumps(result))
cli_utils.CONSOLE.print(table)
[docs]
def model(
ctx: typer.Context,
config: Annotated[
str, typer.Option(*cli_utils.CONFIG_FLAGS, help="Path to the judge config file")
],
inference_config: Annotated[
str,
typer.Option(help="Path to the inference config file"),
],
input_file: Annotated[
Optional[str], typer.Option(help="Path to the input file (jsonl)")
] = None,
output_file: Annotated[
Optional[str], typer.Option(help="Path to the output file (jsonl)")
] = None,
level: cli_utils.LOG_LEVEL_TYPE = None,
):
"""Judge the outputs of a model on a dataset."""
# Delayed imports
from oumi import judge_conversations
from oumi.builders.inference_engines import build_inference_engine
from oumi.core.types.conversation import Conversation
# End imports
judge_extra_args = cli_utils.parse_extra_cli_args(ctx)
config = str(
cli_utils.resolve_and_fetch_config(
config,
)
)
# Load the judge config
judge_config = _load_judge_config(config, judge_extra_args)
# Load the inference config
inference_config = str(
cli_utils.resolve_and_fetch_config(
inference_config,
)
)
inference_extra_args = cli_utils.parse_extra_cli_args(ctx)
model_inference_config: InferenceConfig = InferenceConfig.from_yaml_and_arg_list(
inference_config, inference_extra_args
)
if not model_inference_config.engine:
typer.echo("Inference engine is required.")
raise typer.Exit(code=1)
# Load the dataset
if not input_file:
typer.echo("Input file is required.")
raise typer.Exit(code=1)
input_data = load_jsonlines(input_file)
input_conversations = [Conversation.from_dict(output) for output in input_data]
# Run inference
inference_engine = build_inference_engine(
model_inference_config.engine,
model_params=model_inference_config.model,
remote_params=model_inference_config.remote_params,
generation_params=model_inference_config.generation,
)
model_outputs = inference_engine.infer(
input=input_conversations, inference_config=model_inference_config
)
results = judge_conversations(judge_config, judge_inputs=model_outputs)
if output_file:
typer.echo(f"Saving results to {output_file}")
with jsonlines.open(output_file, mode="w") as writer:
writer.write_all(results)
else:
table = Table(
title="Judge Results",
title_style="bold magenta",
show_edge=False,
show_lines=True,
)
table.add_column("Judgements", style="green")
for result in results:
table.add_row(json.dumps(result))
cli_utils.CONSOLE.print(table)
