# 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 logging
import re
from dataclasses import dataclass, field
from enum import Enum
from oumi.core.configs.params.base_params import BaseParams
from oumi.core.configs.params.judge_params import JudgeOutputType
from oumi.exceptions import OumiConfigError
logger = logging.getLogger(__name__)
# Suffix appended to a criterion id to name its per-criterion explanation field.
EXPLANATION_SUFFIX = "_explanation"
# Criterion ids become XML tag names and JSON object keys. The judge's XML parser
# only recognizes tags matching `\w+`, and a tag name cannot start with a digit, so
# ids are restricted to an identifier-like shape.
_VALID_CRITERION_ID_PATTERN = re.compile(r"^[^\W\d]\w*$")
[docs]
class JudgeAggregation(str, Enum):
"""How to combine per-criterion scores into a single overall score."""
NONE = "none"
"""Do not aggregate; the overall score is left undefined."""
WEIGHTED_MEAN = "weighted_mean"
"""Weighted average of the per-criterion scores, using each criterion's weight."""
MIN = "min"
"""The lowest per-criterion score (a single failing criterion drags it down)."""
ALL = "all"
"""1.0 if every criterion scores 1.0, else 0.0 (the checklist case)."""
[docs]
@dataclass
class JudgeCriterion(BaseParams):
"""A single criterion that a rubric judge evaluates.
Each criterion becomes one field in the judge's structured output, optionally
preceded by its own explanation field.
Example:
>>> criterion = JudgeCriterion( # doctest: +SKIP
... id="groundedness",
... description="Every claim is supported by the provided context.",
... judgment_type=JudgeOutputType.BOOL,
... weight=2.0,
... )
"""
id: str
"""Identifier for this criterion, e.g. "groundedness"."""
description: str
"""What the judge should assess for this criterion."""
judgment_type: JudgeOutputType = field(default=JudgeOutputType.BOOL)
"""The type of output that this criterion's judgment should be provided with."""
judgment_scores: dict[str, float | None] | None = field(default=None)
"""For ENUM judgment_type, the mapping from category names to numeric scores.
Example:
{"excellent": 1.0, "good": 0.7, "poor": 0.3, "not_applicable": None}
"""
include_explanation: bool = field(default=True)
"""Whether the judge should explain this criterion before judging it."""
weight: float = field(default=1.0)
"""Relative weight of this criterion under WEIGHTED_MEAN aggregation."""
@property
def is_scoreable(self) -> bool:
"""Whether this criterion can ever contribute a numeric score.
BOOL criteria are scored 1.0/0.0 automatically. Anything else needs a
`judgment_scores` mapping with at least one label that carries a score --
a mapping whose labels all map to None only constrains the allowed values.
"""
if self.judgment_scores:
return any(score is not None for score in self.judgment_scores.values())
return self.judgment_type == JudgeOutputType.BOOL
@property
def explanation_id(self) -> str:
"""Output field key holding this criterion's explanation."""
return f"{self.id}{EXPLANATION_SUFFIX}"
[docs]
def __post_init__(self):
"""Validate the criterion after initialization."""
self._validate_params()
def _validate_params(self):
"""Validate the criterion for consistency and completeness.
Raises:
OumiConfigError: If the criterion is invalid.
"""
if not self.id or not self.id.strip():
raise OumiConfigError("Criterion `id` cannot be empty")
if not _VALID_CRITERION_ID_PATTERN.match(self.id):
raise OumiConfigError(
f"Criterion id '{self.id}' is invalid: an id must start with a "
"letter or underscore and contain only letters, digits, and "
"underscores, since it is used verbatim as an XML tag name and a "
"JSON key in the judge's response."
)
if self.id.endswith(EXPLANATION_SUFFIX) or self.id == "explanation":
raise OumiConfigError(
f"Criterion id '{self.id}' is reserved: ids named 'explanation' "
f"or ending in '{EXPLANATION_SUFFIX}' name the explanation "
"fields generated by `include_explanation`. Please rename the "
"criterion."
)
if not self.description or not self.description.strip():
raise OumiConfigError(
f"Criterion '{self.id}' must have a non-empty `description`"
)
if self.judgment_type == JudgeOutputType.ENUM and not self.judgment_scores:
raise OumiConfigError(
f"Criterion '{self.id}': judgment_scores must be provided for ENUM "
"judgment_type"
)
if self.judgment_scores is not None:
if not self.judgment_scores:
raise OumiConfigError(
f"Criterion '{self.id}': judgment_scores cannot be empty when "
"provided"
)
# None means "this label carries no score"; anything else must be numeric.
if not all(
score is None or isinstance(score, int | float)
for score in self.judgment_scores.values()
):
raise OumiConfigError(
f"Criterion '{self.id}': judgment_scores values must be numeric, "
"or None for a label that carries no score"
)
# Only ENUM keeps the chosen label as its value. For the other types the
# label is parsed into a bool/int/float, so an unscored label would come
# back as None -- indistinguishable from a failed parse.
if self.judgment_type != JudgeOutputType.ENUM and any(
score is None for score in self.judgment_scores.values()
):
unscored = sorted(
label
for label, score in self.judgment_scores.items()
if score is None
)
raise OumiConfigError(
f"Criterion '{self.id}': judgment_scores labels {unscored} map to "
f"None, which is only supported for ENUM judgment_type (this "
f"criterion is {self.judgment_type.value.upper()}). The value "
"is parsed from the label, so an unscored label could not be "
"told apart from a failed parse. Use judgment_type: ENUM."
)
if self.weight < 0:
raise OumiConfigError(
f"Criterion '{self.id}': weight must be non-negative, got {self.weight}"
)
[docs]
@dataclass
class RubricJudgeParams(BaseParams):
"""Parameters for a rubric judge, which scores several criteria at once.
A rubric judge issues a single inference call per input and asks the judge model
to produce one judgment per criterion, unlike a simple judge which produces one
judgment overall.
Example:
>>> rubric_params = RubricJudgeParams( # doctest: +SKIP
... criteria=[
... JudgeCriterion(
... id="groundedness",
... description="Every claim is supported by the context.",
... include_explanation=True,
... ),
... JudgeCriterion(
... id="formality",
... description=(
... "How formal the response's register is. 'formal' for "
... "professional prose, 'casual' for slang or contractions."
... ),
... judgment_type=JudgeOutputType.ENUM,
... judgment_scores={"formal": 1.0, "neutral": 0.5, "casual": 0.0},
... ),
... ],
... aggregation=JudgeAggregation.WEIGHTED_MEAN,
... )
"""
criteria: list[JudgeCriterion] = field(default_factory=list)
"""The criteria to evaluate, in the order the judge should address them."""
aggregation: JudgeAggregation = field(default=JudgeAggregation.WEIGHTED_MEAN)
"""How to combine per-criterion scores into the output's `aggregate_score`."""
[docs]
def __post_init__(self):
"""Validate the parameters after initialization."""
self._validate_params()
self._warn_about_aggregation_gaps()
def _warn_about_aggregation_gaps(self) -> None:
"""Warn about config that silently has no effect on the aggregate score."""
self._warn_about_ignored_weights()
self._warn_about_unscored_criteria()
def _warn_about_ignored_weights(self) -> None:
"""Warn when weights are set under an aggregation that does not read them.
Only WEIGHTED_MEAN consults `weight`; MIN and ALL are order statistics and
NONE aggregates nothing, so a weight set under those is a silent no-op.
"""
if self.aggregation == JudgeAggregation.WEIGHTED_MEAN:
return
weighted = [c.id for c in self.criteria if c.weight != 1.0]
if weighted:
logger.warning(
f"Criteria {sorted(weighted)} set a non-default `weight`, but "
f"'{self.aggregation.value}' aggregation ignores weights. The weights "
"have no effect. Use 'weighted_mean' aggregation, or remove them."
)
def _warn_about_unscored_criteria(self) -> None:
"""Warn about criteria that cannot contribute to the aggregate score."""
if self.aggregation == JudgeAggregation.NONE:
return
unscored = [c.id for c in self.criteria if not c.is_scoreable]
if unscored:
logger.warning(
f"Criteria {sorted(unscored)} produce no numeric score and will be "
f"excluded from the '{self.aggregation.value}' aggregate score. "
"Add `judgment_scores` to include them."
)
def _validate_params(self):
"""Validate the parameters for consistency and completeness.
Raises:
OumiConfigError: If parameters are invalid.
"""
if not self.criteria:
raise OumiConfigError("`criteria` cannot be empty for a rubric judge")
# Criterion ids must be unique. They cannot collide with a generated
# explanation field, because JudgeCriterion reserves that suffix.
seen_ids: set[str] = set()
for criterion in self.criteria:
if criterion.id in seen_ids:
raise OumiConfigError(
f"Duplicate criterion id: '{criterion.id}'. Criterion ids must "
"be unique."
)
seen_ids.add(criterion.id)
if self.aggregation == JudgeAggregation.WEIGHTED_MEAN and not any(
criterion.weight > 0 for criterion in self.criteria
):
raise OumiConfigError(
"WEIGHTED_MEAN aggregation requires at least one criterion with a "
"positive weight; all weights are zero."
)
