Add LLM-as-a-Judge evaluator using MLflow and OpenAI-compatible endpoints

ai4rag/components/optimization/__init__.py

@@ -2,7 +2,10 @@
 # Copyright IBM Corp. 2026
 # SPDX-License-Identifier: Apache-2.0
 # -----------------------------------------------------------------------------
-from ai4rag.components.optimization.rag_templates_optimization import OptimizationResult, run_rag_optimization
+from ai4rag.components.optimization.rag_templates_optimization import (
+    OptimizationResult,
+    run_rag_optimization,
+)
 from ai4rag.components.optimization.search_space_preparation import SearchSpaceReport, prepare_search_space_report
 
 __all__ = [

ai4rag/components/optimization/judge_selection.py

@@ -0,0 +1,165 @@
+# -----------------------------------------------------------------------------
+# Copyright IBM Corp. 2026
+# SPDX-License-Identifier: Apache-2.0
+# -----------------------------------------------------------------------------
+import logging
+from typing import Any
+
+import numpy as np
+from docling_core.types.doc import DoclingDocument
+
+from ai4rag import handler
+from ai4rag.core.experiment.benchmark_data import BenchmarkData
+from ai4rag.core.experiment.utils import build_evaluation_data, query_rag
+from ai4rag.evaluator.base_evaluator import MetricType
+from ai4rag.evaluator.llmaj_evaluator import LLMaJConfig, LLMaJEvaluator
+from ai4rag.evaluator.unitxt_evaluator import UnitxtEvaluator
+from ai4rag.rag.chunking.langchain_chunker import LangChainChunker
+from ai4rag.rag.embedding.base_model import BaseEmbeddingModel
+from ai4rag.rag.foundation_models.base_model import BaseFoundationModel
+from ai4rag.rag.retrieval.retriever import Retriever
+from ai4rag.rag.template.simple_rag_template import SimpleRAG
+from ai4rag.rag.vector_store.chroma import ChromaVectorStore
+
+_logger = logging.getLogger("judge-selection")
+_logger.addHandler(handler)
+
+
+def calibration_subset_size(total_rows: int) -> int:
+    """Return calibration row count: min(20, 10% of benchmark rows)."""
+    if total_rows <= 0:
+        return 0
+    return max(1, min(20, int(total_rows * 0.1)))
+
+
+def select_judge_model(
+    *,
+    evaluator: str,
+    judge_model_id: str | None,
+    generation_models: list[BaseFoundationModel],
+    embedding_models: list[BaseEmbeddingModel],
+    benchmark_data: BenchmarkData,
+    documents: list[DoclingDocument],
+    ogx_base_url: str,
+    ogx_api_key: str,
+    max_threads: int = 10,
+) -> str | None:
+    """
+    Resolve the judge model for an optimization run.
+
+    When ``evaluator`` is ``unitxt``, returns ``None``. When ``judge_model_id``
+    is provided, returns it unchanged. Otherwise auto-selects from the generation
+    model pool using faithfulness calibration against Unitxt.
+    """
+    if evaluator != "judge":
+        return None
+
+    if judge_model_id:
+        return judge_model_id
+
+    candidates = [model.model_id for model in generation_models]
+    if not candidates:
+        raise ValueError("At least one generation model is required to select a judge model.")
+    if len(candidates) == 1:
+        return candidates[0]
+
+    return _calibrate_judge_model(
+        candidates=candidates,
+        generation_models=generation_models,
+        embedding_models=embedding_models,
+        benchmark_data=benchmark_data,
+        documents=documents,
+        ogx_base_url=ogx_base_url,
+        ogx_api_key=ogx_api_key,
+        max_threads=max_threads,
+    )
+
+
+def _calibrate_judge_model(
+    *,
+    candidates: list[str],
+    generation_models: list[BaseFoundationModel],
+    embedding_models: list[BaseEmbeddingModel],
+    benchmark_data: BenchmarkData,
+    documents: list[DoclingDocument],
+    ogx_base_url: str,
+    ogx_api_key: str,
+    max_threads: int,
+) -> str:
+    subset_size = calibration_subset_size(len(benchmark_data.questions))
+    calibration_benchmark = benchmark_data.get_random_sample(n_records=subset_size, random_seed=17)
+    eval_data = _run_reference_rag(
+        foundation_model=generation_models[0],
+        embedding_model=embedding_models[0],
+        documents=documents,
+        benchmark_data=calibration_benchmark,
+        max_threads=max_threads,
+    )
+
+    unitxt_scores = UnitxtEvaluator().evaluate_metrics(eval_data, [MetricType.FAITHFULNESS])
+    reference = _ordered_question_scores(unitxt_scores, MetricType.FAITHFULNESS)
+
+    rankings: list[dict[str, Any]] = []
+    for model_id in candidates:
+        judge = LLMaJEvaluator(
+            LLMaJConfig(
+                base_url=ogx_base_url.rstrip("/") + "/v1",
+                api_key=ogx_api_key,
+                model=model_id,
+            )
+        )
+        judge_scores = judge.evaluate_metrics(eval_data, [MetricType.FAITHFULNESS])
+        candidate_scores = _ordered_question_scores(judge_scores, MetricType.FAITHFULNESS)
+        correlation = _pearson_correlation(reference, candidate_scores)
+        rankings.append({"model_id": model_id, "judge_calibration_score": correlation})
+        _logger.info("Judge calibration for %s: correlation=%.4f", model_id, correlation)
+
+    rankings.sort(
+        key=lambda item: (
+            item["judge_calibration_score"] if item["judge_calibration_score"] is not None else -2.0,
+            -candidates.index(item["model_id"]),
+            item["model_id"],
+        ),
+        reverse=True,
+    )
+    selected = rankings[0]["model_id"]
+    _logger.info("Selected judge model: %s (calibration score=%s)", selected, rankings[0]["judge_calibration_score"])
+    return selected
+
+
+def _run_reference_rag(
+    *,
+    foundation_model: BaseFoundationModel,
+    embedding_model: BaseEmbeddingModel,
+    documents: list[DoclingDocument],
+    benchmark_data: BenchmarkData,
+    max_threads: int,
+) -> list:
+    chunker = LangChainChunker(chunk_size=512, method="recursive", chunk_overlap=128)
+    chunks = chunker.split_documents(documents)
+    vector_store = ChromaVectorStore(embedding_model=embedding_model, collection_name="judge_calibration")
+    vector_store.add_documents(chunks)
+    retriever = Retriever(vector_store=vector_store, number_of_chunks=3, method="simple", search_mode="vector")
+    rag = SimpleRAG(foundation_model=foundation_model, retriever=retriever)
+    inference_response = query_rag(
+        rag=rag,
+        questions=list(benchmark_data.questions),
+        max_threads=max_threads,
+    )
+    return build_evaluation_data(benchmark_data=benchmark_data, inference_response=inference_response)
+
+
+def _ordered_question_scores(evaluation_result: dict, metric: str) -> list[float | None]:
+    question_scores = (evaluation_result.get("question_scores") or {}).get(metric) or {}
+    return [question_scores.get(qid) for qid in sorted(question_scores.keys())]
+
+
+def _pearson_correlation(reference: list[float | None], candidate: list[float | None]) -> float | None:
+    paired = [(a, b) for a, b in zip(reference, candidate) if a is not None and b is not None]
+    if len(paired) < 2:
+        return None
+    ref_vals = np.array([pair[0] for pair in paired], dtype=float)
+    cand_vals = np.array([pair[1] for pair in paired], dtype=float)
+    if np.std(ref_vals) == 0 or np.std(cand_vals) == 0:
+        return None
+    return float(np.corrcoef(ref_vals, cand_vals)[0, 1])

ai4rag/components/optimization/rag_templates_optimization.py

@@ -18,6 +18,9 @@
 from ai4rag.components.utils.docling_io import load_docling_documents
 from ai4rag.core.experiment.experiment import AI4RAGExperiment
 from ai4rag.core.hpo.gam_opt import GAMOptSettings
+from ai4rag.evaluator.base_evaluator import MetricType
+from ai4rag.evaluator.llmaj_evaluator import LLMaJConfig, LLMaJEvaluator
+from ai4rag.evaluator.unitxt_evaluator import UnitxtEvaluator
 from ai4rag.rag.embedding.ogx import OGXEmbeddingModel
 from ai4rag.rag.foundation_models.base_model import Language
 from ai4rag.rag.foundation_models.ogx import OGXFoundationModel
@@ -32,6 +35,11 @@
 MIN_MAX_RAG_PATTERNS_RANGE = (4, 20)
 DEFAULT_METRIC = "faithfulness"
 SUPPORTED_OPTIMIZATION_METRICS = frozenset({"faithfulness", "answer_correctness", "context_correctness"})
+STANDARD_EVALUATION_METRICS = (
+    MetricType.ANSWER_CORRECTNESS,
+    MetricType.FAITHFULNESS,
+    MetricType.CONTEXT_CORRECTNESS,
+)
 
 
 @dataclass
@@ -61,6 +69,8 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
     input_data_key: str = "",
     optimization_settings: dict | None = None,
     max_threads: int = 10,
+    evaluator: str = "judge",
+    judge_model_id: str | None = None,
 ) -> OptimizationResult:
     """Run a full AI4RAG optimization experiment and generate output artefacts.
 
@@ -96,6 +106,11 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
         RAG service during benchmark evaluation.  Lower values reduce
         per-request concurrency (useful when each request carries more
         retrieved context).  Defaults to ``10``.
+    evaluator
+        Evaluation backend: ``"judge"`` (default) or ``"unitxt"``.
+    judge_model_id
+        Judge model identifier for ``evaluator="judge"``. When omitted, the
+        value from the search-space report is used.
 
     Returns
     -------
@@ -120,21 +135,63 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
         raise ValueError("vector_io_provider_id must be a non-empty string.")
     vector_io_provider_id = vector_io_provider_id.strip()
 
+    if evaluator not in ("judge", "unitxt"):
+        raise ValueError(f"Evaluator {evaluator!r} is not supported. Supported evaluators are ['judge', 'unitxt'].")
+
     settings = _validate_optimization_settings(optimization_settings)
     optimization_metric = settings.get("metric") or DEFAULT_METRIC
+
     if optimization_metric not in SUPPORTED_OPTIMIZATION_METRICS:
         raise ValueError(
             f"Optimization metric {optimization_metric} is not supported. "
-            f"Select one of {SUPPORTED_OPTIMIZATION_METRICS}."
+            f"Select one of {sorted(SUPPORTED_OPTIMIZATION_METRICS)}."
+        )
+
+    if evaluator == "judge" and not judge_model_id:
+        raise ValueError(
+            "judge_model_id is required when evaluator='judge'. "
+            "Provide it explicitly or run search-space preparation first."
         )
 
     documents = load_docling_documents(extracted_text_path)
 
     with open(search_space_report_path, "r", encoding="utf-8") as f:
         search_space_raw: dict[str, Any] = json.load(f)
 
+    evaluation_block = search_space_raw.get("evaluation") or {}
+    resolved_evaluator = evaluator
+    resolved_judge_model_id = judge_model_id or evaluation_block.get("judge_model_id")
+    if resolved_evaluator == "judge" and not resolved_judge_model_id:
+        raise ValueError(
+            "judge_model_id is required when evaluator='judge'. "
+            "Provide it explicitly or run search-space preparation first."
+        )
+
+    evaluator_kwargs: dict[str, Any]
+    if resolved_evaluator == "unitxt":
+        evaluator_kwargs = {
+            "evaluator": UnitxtEvaluator(),
+            "metrics": STANDARD_EVALUATION_METRICS,
+        }
+    else:
+        config = LLMaJConfig(
+            base_url=ogx_client.base_url.rstrip("/") + "/v1",
+            api_key=getattr(ogx_client, "api_key", "") or "",
+            model=resolved_judge_model_id or "",
+        )
+        evaluator_kwargs = {
+            "evaluator": LLMaJEvaluator(config),
+            "metrics": STANDARD_EVALUATION_METRICS,
+        }
+
+    evaluation_config = {"evaluator": resolved_evaluator}
+    if resolved_evaluator == "judge":
+        evaluation_config["judge_model_id"] = resolved_judge_model_id
+
     params: list[Parameter] = []
     for param_name, values in search_space_raw.items():
+        if param_name in ("evaluation",):
+            continue
         if param_name in ("foundation_model", "embedding_model"):
             values = [_deserialize_model(m, ogx_client) for m in values]
         params.append(Parameter(param_name, "C", values=values))
@@ -161,7 +218,9 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
         documents=documents,
         optimization_metric=optimization_metric,
         ogx_vector_io_provider_id=vector_io_provider_id,
-        max_threads=max_threads,
+        inference_max_threads=max_threads,
+        evaluation_config=evaluation_config,
+        **evaluator_kwargs,
     )
 
     # --- Run the optimization loop ---
@@ -199,12 +258,9 @@ def run_rag_optimization(  # pylint: disable=too-many-locals,too-many-arguments,
             ogx_base_url=ogx_base_url,
         )
 
-        # Attach scores to pattern data and write pattern.json
-
         with (patt_dir / "pattern.json").open("w+", encoding="utf-8") as f:
             json_dump(pattern_data, f, indent=2)
 
-        # Write evaluation results
         evaluation_result_list = pattern.get("evaluation_results", [])
         with (patt_dir / "evaluation_results.json").open("w+", encoding="utf-8") as f:
             json_dump(evaluation_result_list, f, indent=2)

ai4rag/components/optimization/search_space_preparation.py

@@ -12,6 +12,7 @@
 from ogx_client import OgxClient
 
 from ai4rag import handler
+from ai4rag.components.optimization.judge_selection import select_judge_model
 from ai4rag.components.utils.docling_io import load_docling_documents
 from ai4rag.core.experiment.benchmark_data import BenchmarkData
 from ai4rag.core.experiment.mps import ModelsPreSelector
@@ -114,6 +115,8 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     random_seed: int = _DEFAULT_SEED,
     chunking_methods: list[str] | None = None,
     inference_max_threads: int = 10,
+    evaluator: str = "judge",
+    judge_model_id: str | None = None,
 ) -> SearchSpaceReport:
     """Run model pre-selection and prepare a search-space report.
 
@@ -158,6 +161,11 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         RAG service during benchmark evaluation.  Lower values reduce
         per-request concurrency (useful when each request carries more
         retrieved context).  Defaults to ``10``.
+    evaluator
+        Evaluation backend: ``"judge"`` (default) or ``"unitxt"``.
+    judge_model_id
+        Optional judge model identifier. When omitted and ``evaluator`` is
+        ``"judge"``, a judge model is auto-selected during preparation.
 
     Returns
     -------
@@ -175,6 +183,9 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
     if metric not in SUPPORTED_METRICS:
         raise ValueError(f"Metric {metric!r} is not supported. Supported metrics are {list(SUPPORTED_METRICS)}.")
 
+    if evaluator not in ("judge", "unitxt"):
+        raise ValueError(f"Evaluator {evaluator!r} is not supported. Supported evaluators are ['judge', 'unitxt'].")
+
     _validate_model_list(embedding_models, "embedding_models")
     _validate_model_list(generation_models, "generation_models")
     _validate_chunking_methods(chunking_methods)
@@ -240,6 +251,25 @@ def prepare_search_space_report(  # pylint: disable=too-many-locals,too-many-arg
         verbose_repr["chunking_method"] = chunking_methods
         _logger.info("Chunking methods constrained to: %s", verbose_repr["chunking_method"])
 
+    evaluation_block: dict[str, str] = {"evaluator": evaluator}
+    if evaluator == "judge":
+        ogx_api_key = getattr(ogx_client, "api_key", "") or ""
+        resolved_judge_id = select_judge_model(
+            evaluator=evaluator,
+            judge_model_id=judge_model_id,
+            generation_models=selected_models["foundation_model"],
+            embedding_models=selected_models["embedding_model"],
+            benchmark_data=benchmark_data,
+            documents=documents,
+            ogx_base_url=ogx_client.base_url,
+            ogx_api_key=ogx_api_key,
+            max_threads=inference_max_threads,
+        )
+        if not resolved_judge_id:
+            raise ValueError("Failed to resolve judge_model_id for evaluator='judge'.")
+        evaluation_block["judge_model_id"] = resolved_judge_id
+    verbose_repr["evaluation"] = evaluation_block
+
     return SearchSpaceReport(
         search_space=verbose_repr,
         selected_models=selected_models,