"""Pipeline predictor - Handles prediction mode execution.
This module provides the Predictor class for running predictions using trained
pipelines on new datasets.
Supports two prediction paths:
- Bundle path: ``.n4a`` bundles loaded via BundleLoader (self-contained).
- Resolver path: Prediction dicts / folders resolved via PredictionResolver,
with optional minimal pipeline extraction from execution traces.
"""
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional, Union
import numpy as np
from nirs4all.core.logging import Nirs4allLogger, get_logger
if TYPE_CHECKING:
from nirs4all.pipeline.resolver import ResolvedPrediction
from nirs4all.pipeline.runner import PipelineRunner
from nirs4all.pipeline.trace import ExecutionTrace, MinimalPipeline
from nirs4all.data.config import DatasetConfigs
from nirs4all.data.dataset import SpectroDataset
from nirs4all.data.predictions import Predictions
from nirs4all.pipeline.config.context import (
DataSelector,
ExecutionContext,
LoaderArtifactProvider,
PipelineState,
RuntimeContext,
StepMetadata,
)
from nirs4all.pipeline.execution.builder import ExecutorBuilder
from nirs4all.pipeline.storage.artifacts.artifact_loader import ArtifactLoader
logger: Nirs4allLogger = get_logger(__name__)
@dataclass(frozen=True)
class _PredictionStrategy:
"""Plan for executing one prediction replay strategy."""
name: str
mode: str # "minimal" or "full"
steps: list[Any]
minimal_pipeline: Any | None
artifact_provider_factory: Callable[[], Any]
[docs]
class Predictor:
"""Handles prediction using trained pipelines.
This class manages the prediction workflow: loading saved models,
replaying pipeline configurations, and generating predictions on new data.
When ``use_minimal_pipeline=True`` (default), the predictor will:
1. Check if an execution trace is available for the prediction.
2. Extract the minimal pipeline (only required steps) from the trace.
3. Execute only those steps, significantly reducing prediction time.
Attributes:
runner: Parent PipelineRunner instance.
pipeline_uid: Unique identifier for the pipeline.
artifact_loader: Loader for trained model artifacts.
config_path: Path to the pipeline configuration.
target_model: Metadata for the target model.
use_minimal_pipeline: Whether to use minimal pipeline execution.
"""
def __init__(self, runner: "PipelineRunner", use_minimal_pipeline: bool = True):
"""Initialize predictor.
Args:
runner: Parent PipelineRunner instance.
use_minimal_pipeline: If True, use minimal pipeline execution when
execution traces are available.
"""
self.runner = runner
self.pipeline_uid: str | None = None
self.artifact_loader: ArtifactLoader | None = None
self.config_path: str | None = None
self.target_model: dict[str, Any] | None = None
self.use_minimal_pipeline = use_minimal_pipeline
self._execution_trace: ExecutionTrace | None = None
self._minimal_pipeline: MinimalPipeline | None = None
self._resolved: ResolvedPrediction | None = None
[docs]
def predict(
self,
prediction_obj: dict[str, Any] | str,
dataset: DatasetConfigs | SpectroDataset | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
dataset_name: str = "prediction_dataset",
all_predictions: bool = False,
verbose: int = 0,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Run prediction using a saved model on new dataset.
Args:
prediction_obj: Model identifier (dict with config_path or prediction ID).
dataset: New dataset to predict on.
dataset_name: Name for the dataset.
all_predictions: If True, return all predictions; if False, return single best.
verbose: Verbosity level.
Returns:
If ``all_predictions=False``: ``(y_pred, predictions)``
If ``all_predictions=True``: ``(predictions_dict, predictions)``
"""
logger.starting("Starting Nirs4all prediction(s)")
# Handle bundle files (.n4a) directly via BundleLoader
if isinstance(prediction_obj, str) and (prediction_obj.endswith(".n4a") or prediction_obj.endswith(".n4a.py")):
return self._predict_from_bundle(prediction_obj, dataset, dataset_name, all_predictions, verbose)
# Normalize dataset
dataset_config = self.runner.orchestrator._normalize_dataset(dataset, dataset_name)
# Setup prediction mode
self.runner.mode = "predict"
self.runner.verbose = verbose
# Resolve prediction source via PredictionResolver
self._resolve_prediction(prediction_obj, verbose)
# Try minimal pipeline execution if trace is available
if self.use_minimal_pipeline and self._execution_trace is not None:
return self._predict_with_minimal_pipeline(dataset_config, all_predictions, verbose)
# Fallback: Full pipeline execution
return self._predict_full_pipeline(dataset_config, all_predictions, verbose)
# -----------------------------------------------------------------
# Resolution
# -----------------------------------------------------------------
def _resolve_prediction(self, prediction_obj: dict[str, Any] | str, verbose: int) -> None:
"""Resolve a prediction object to executable components.
Extracts ``config_path``, ``target_model``, ``artifact_loader`` and
optionally an ``_execution_trace`` from the prediction source.
Args:
prediction_obj: Prediction dict or string identifier.
verbose: Verbosity level.
"""
from nirs4all.pipeline.resolver import PredictionResolver
resolver = PredictionResolver(
workspace_path=self.runner.workspace_path,
runs_dir=self.runner.workspace_path,
store=self.runner.store,
)
resolved = resolver.resolve(prediction_obj, verbose=verbose)
self._resolved = resolved
# Extract target model metadata
target_model = {k: v for k, v in prediction_obj.items() if k not in ("y_pred", "y_true", "X")} if isinstance(prediction_obj, dict) else dict(resolved.target_model) if resolved.target_model else {}
if not target_model or "model_name" not in target_model:
raise ValueError(
"No model information found in prediction. "
"Please provide a prediction dict with 'model_name' and 'pipeline_uid' keys."
)
self.target_model = target_model
self.runner.target_model = target_model
self.pipeline_uid = resolved.pipeline_uid
self.config_path = resolved.pipeline_uid
self.artifact_loader = None
# Create artifact loader from resolved artifact provider
if resolved.artifact_provider and hasattr(resolved.artifact_provider, "artifact_loader"):
self.artifact_loader = resolved.artifact_provider.artifact_loader
# Cache execution trace
self._execution_trace = resolved.trace
# -----------------------------------------------------------------
# Minimal pipeline path
# -----------------------------------------------------------------
def _predict_with_minimal_pipeline(
self,
dataset_config: DatasetConfigs,
all_predictions: bool,
verbose: int,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Execute prediction using minimal pipeline extracted from trace.
Args:
dataset_config: Dataset configuration.
all_predictions: Whether to return all predictions.
verbose: Verbosity level.
Returns:
Same as ``predict()`` method.
"""
logger.info("Using minimal pipeline execution")
strategy = self._make_minimal_strategy()
return self._run_prediction_strategy(
strategy=strategy,
dataset_config=dataset_config,
all_predictions=all_predictions,
verbose=verbose,
)
# -----------------------------------------------------------------
# Full pipeline path
# -----------------------------------------------------------------
def _predict_full_pipeline(
self,
dataset_config: DatasetConfigs,
all_predictions: bool,
verbose: int,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Execute prediction using full pipeline replay.
Args:
dataset_config: Dataset configuration.
all_predictions: Whether to return all predictions.
verbose: Verbosity level.
Returns:
Same as ``predict()`` method.
"""
strategy = self._make_full_strategy()
return self._run_prediction_strategy(
strategy=strategy,
dataset_config=dataset_config,
all_predictions=all_predictions,
verbose=verbose,
)
def _make_minimal_strategy(self) -> _PredictionStrategy:
"""Build the minimal-pipeline replay strategy plan."""
from nirs4all.pipeline.minimal_predictor import MinimalArtifactProvider
from nirs4all.pipeline.trace import TraceBasedExtractor
assert self._resolved is not None
assert self._execution_trace is not None
extractor = TraceBasedExtractor()
full_steps = list(self._resolved.minimal_pipeline)
target_branch_name = self.target_model.get("branch_name") if self.target_model else None
target_branch_id = self.target_model.get("branch_id") if self.target_model else None
target_branch_path = self.target_model.get("branch_path") if self.target_model else None
if target_branch_name:
minimal_pipeline = extractor.extract_for_branch_name(
trace=self._execution_trace,
branch_name=target_branch_name,
full_pipeline=full_steps,
)
elif target_branch_path:
minimal_pipeline = extractor.extract_for_branch(
trace=self._execution_trace,
branch_path=target_branch_path,
full_pipeline=full_steps,
)
elif target_branch_id is not None:
minimal_pipeline = extractor.extract_for_branch(
trace=self._execution_trace,
branch_path=[target_branch_id],
full_pipeline=full_steps,
)
else:
minimal_pipeline = extractor.extract(
trace=self._execution_trace,
full_pipeline=full_steps,
up_to_model=True,
)
self._minimal_pipeline = minimal_pipeline
logger.debug(
f"Minimal pipeline: {minimal_pipeline.get_step_count()} steps "
f"(from {len(full_steps)} total)"
)
target_sub_index = None
target_model_name = None
if self.target_model:
model_artifact_id = self.target_model.get("model_artifact_id")
if model_artifact_id:
target_sub_index = self._get_substep_from_artifact(model_artifact_id)
else:
target_model_name = self.target_model.get("model_name")
def _artifact_provider_factory() -> Any:
return MinimalArtifactProvider(
minimal_pipeline=minimal_pipeline,
artifact_loader=self.artifact_loader,
target_sub_index=target_sub_index,
target_model_name=target_model_name,
)
return _PredictionStrategy(
name="minimal",
mode="minimal",
steps=[step.step_config for step in minimal_pipeline.steps],
minimal_pipeline=minimal_pipeline,
artifact_provider_factory=_artifact_provider_factory,
)
def _make_full_strategy(self) -> _PredictionStrategy:
"""Build the full-pipeline replay strategy plan."""
assert self._resolved is not None
resolved = self._resolved
def _artifact_provider_factory() -> Any:
artifact_provider = resolved.artifact_provider
if artifact_provider is None and self.artifact_loader:
artifact_provider = LoaderArtifactProvider(loader=self.artifact_loader)
return artifact_provider
return _PredictionStrategy(
name="full",
mode="full",
steps=list(resolved.minimal_pipeline),
minimal_pipeline=None,
artifact_provider_factory=_artifact_provider_factory,
)
def _run_prediction_strategy(
self,
strategy: _PredictionStrategy,
dataset_config: DatasetConfigs,
all_predictions: bool,
verbose: int,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Execute prediction replay using a strategy plan."""
run_predictions = Predictions()
for config, name in dataset_config.configs:
dataset_obj = dataset_config.get_dataset(config, name)
config_predictions = Predictions()
selector = self._build_prediction_selector(dataset_obj)
context = self._build_prediction_context(selector)
executor = self._build_prediction_executor(verbose)
runtime_context = self._build_prediction_runtime_context(
executor=executor,
artifact_provider=strategy.artifact_provider_factory(),
)
if strategy.mode == "minimal":
executor.execute_minimal(
steps=strategy.steps,
minimal_pipeline=strategy.minimal_pipeline,
dataset=dataset_obj,
context=context,
runtime_context=runtime_context,
prediction_store=config_predictions,
)
else:
executor.execute(
strategy.steps,
"prediction",
dataset_obj,
context,
runtime_context,
config_predictions,
)
run_predictions.merge_predictions(config_predictions)
return self._process_prediction_results(run_predictions, all_predictions)
def _build_prediction_selector(self, dataset_obj: SpectroDataset) -> DataSelector:
"""Build prediction selector with optional branch filters."""
selector = DataSelector(
partition=None,
processing=[["raw"]] * dataset_obj.features_sources(),
layout="2d",
concat_source=True,
)
target_branch_id = self.target_model.get("branch_id") if self.target_model else None
target_branch_name = self.target_model.get("branch_name") if self.target_model else None
target_branch_path = self.target_model.get("branch_path") if self.target_model else None
if target_branch_id is not None or target_branch_name:
bp = target_branch_path or ([target_branch_id] if target_branch_id is not None else None)
selector = selector.with_branch(
branch_id=target_branch_id,
branch_name=target_branch_name,
branch_path=bp,
)
return selector
@staticmethod
def _build_prediction_context(selector: DataSelector) -> ExecutionContext:
"""Build execution context for prediction replay."""
return ExecutionContext(
selector=selector,
state=PipelineState(y_processing="numeric", step_number=0, mode="predict"),
metadata=StepMetadata(),
)
def _build_prediction_executor(self, verbose: int) -> Any:
"""Build configured executor used by both replay strategies."""
return (
ExecutorBuilder()
.with_workspace(self.runner.workspace_path)
.with_verbose(verbose)
.with_mode("predict")
.with_save_artifacts(False)
.with_save_charts(self.runner.save_charts)
.with_continue_on_error(self.runner.continue_on_error)
.with_show_spinner(self.runner.show_spinner)
.with_plots_visible(self.runner.plots_visible)
.with_figure_refs(self.runner._figure_refs)
.with_artifact_loader(self.artifact_loader)
.build()
)
def _build_prediction_runtime_context(self, executor: Any, artifact_provider: Any) -> RuntimeContext:
"""Build runtime context shared by full and minimal prediction replay."""
return RuntimeContext(
artifact_loader=self.artifact_loader,
artifact_provider=artifact_provider,
step_runner=executor.step_runner,
target_model=self.target_model,
explainer=self.runner.explainer,
)
# -----------------------------------------------------------------
# Result processing
# -----------------------------------------------------------------
def _process_prediction_results(
self,
run_predictions: Predictions,
all_predictions: bool,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Process prediction results and return in requested format.
Args:
run_predictions: Predictions object with results.
all_predictions: Whether to return all predictions.
Returns:
Formatted prediction results.
"""
if all_predictions:
res: dict[str, Any] = {}
for pred in run_predictions.to_dicts():
ds = pred["dataset_name"]
if ds not in res:
res[ds] = {}
res[ds][pred["id"]] = pred["y_pred"]
return res, run_predictions
# Get single prediction matching target model
tm = self.target_model or {}
target_fold_id = tm.get("fold_id", None)
is_aggregated_fold = target_fold_id in ("avg", "w_avg")
is_refit_fold = target_fold_id == "final"
filter_kwargs: dict[str, Any] = {
"model_name": tm.get("model_name"),
"step_idx": tm.get("step_idx"),
}
if not is_aggregated_fold and not is_refit_fold:
filter_kwargs["fold_id"] = target_fold_id
target_branch_id = tm.get("branch_id")
if target_branch_id is not None:
filter_kwargs["branch_id"] = target_branch_id
candidates = run_predictions.filter_predictions(**filter_kwargs)
non_empty = [p for p in candidates if len(p["y_pred"]) > 0]
single_pred = self._aggregate_fold_predictions(non_empty, str(target_fold_id)) if is_aggregated_fold and len(non_empty) > 1 else non_empty[0] if non_empty else (candidates[0] if candidates else None)
if single_pred is None:
raise ValueError("No matching prediction found for the specified model criteria.")
logger.success(f"Predicted with: {single_pred['model_name']} [{single_pred['id']}]")
return single_pred["y_pred"], run_predictions
def _aggregate_fold_predictions(
self,
non_empty: list[dict[str, Any]],
target_fold_id: str,
) -> dict[str, Any] | None:
"""Aggregate predictions across folds.
Args:
non_empty: Non-empty prediction entries.
target_fold_id: Target fold id ('avg' or 'w_avg').
Returns:
Aggregated prediction dict, or first candidate.
"""
tm = self.target_model or {}
fold_weights = tm.get("weights")
fold_preds: dict[Any, np.ndarray] = {}
for p in non_empty:
fid = p.get("fold_id")
if fid not in ("avg", "w_avg") and fid is not None and fid not in fold_preds:
fold_preds[fid] = p["y_pred"]
if not fold_preds:
return non_empty[0] if non_empty else None
sorted_folds = sorted(fold_preds.keys())
y_arrays = [fold_preds[fid] for fid in sorted_folds]
has_valid_weights = (
target_fold_id == "w_avg"
and fold_weights is not None
and len(fold_weights) >= len(sorted_folds)
)
if has_valid_weights and fold_weights is not None:
wl = [fold_weights.get(fid, 1.0) for fid in sorted_folds] if isinstance(fold_weights, dict) else [fold_weights[i] for i in range(len(sorted_folds))]
total_w = sum(wl)
y_pred = sum(w * np.array(y) for w, y in zip(wl, y_arrays, strict=False)) / total_w
else:
y_pred = np.mean(y_arrays, axis=0)
single_pred = non_empty[0].copy()
single_pred["y_pred"] = y_pred
single_pred["fold_id"] = target_fold_id
return single_pred
# -----------------------------------------------------------------
# Bundle path
# -----------------------------------------------------------------
def _predict_from_bundle(
self,
bundle_path: str,
dataset: DatasetConfigs | SpectroDataset | np.ndarray | tuple[np.ndarray, ...] | dict | list[dict] | str | list[str],
dataset_name: str,
all_predictions: bool,
verbose: int,
) -> tuple[np.ndarray, Predictions] | tuple[dict[str, Any], Predictions]:
"""Predict from an exported bundle file.
Args:
bundle_path: Path to ``.n4a`` bundle file.
dataset: Dataset to predict on.
dataset_name: Name for the dataset.
all_predictions: Whether to return all predictions.
verbose: Verbosity level.
Returns:
Predictions.
"""
from nirs4all.pipeline.bundle import BundleLoader
logger.debug(f"Loading bundle: {bundle_path}")
loader = BundleLoader(bundle_path)
from nirs4all.data.raw_multisource import RawMultiSourceDataset
prediction_metadata: dict[str, np.ndarray] = {}
native_routing = loader.get_partitioner_routing() or {}
native_metadata = any(info.get("native_replay") == "by_metadata_concat" for info in native_routing.values())
if isinstance(dataset, RawMultiSourceDataset):
if not loader.relation_replay_manifest:
raise ValueError("RawMultiSourceDataset bundle prediction requires a relation replay manifest.")
X_data: Any = dataset
elif isinstance(dataset, dict) and "X" in dataset:
X_data = np.asarray(dataset["X"])
supplied_metadata = dataset.get("metadata")
if isinstance(supplied_metadata, dict):
prediction_metadata = {str(key): np.asarray(value) for key, value in supplied_metadata.items()}
elif native_metadata and isinstance(dataset, np.ndarray):
required = loader.get_required_metadata_columns()
raise ValueError(f"by_metadata bundle prediction requires metadata column(s) {required!r} alongside X")
else:
dataset_config = self.runner.orchestrator._normalize_dataset(dataset, dataset_name)
X_data = None
for data_config, name in dataset_config.configs:
dataset_obj = dataset_config.get_dataset(data_config, name)
raw = dataset_obj.x({})
X_data = np.concatenate(raw, axis=1) if isinstance(raw, list) else raw
if native_metadata:
for column in loader.get_required_metadata_columns():
try:
prediction_metadata[column] = np.asarray(dataset_obj.metadata_column(column, {}))
except (KeyError, ValueError) as exc:
raise ValueError(f"by_metadata bundle prediction requires metadata column {column!r}") from exc
break
if X_data is None:
raise ValueError("No data found in dataset for prediction")
y_pred = loader.predict_with_metadata(X_data, prediction_metadata) if native_metadata else loader.predict(X_data)
model_name = "bundle_model"
if loader.metadata:
model_name = loader.metadata.original_manifest.get("name", "bundle_model")
metadata: dict[str, Any] = {}
if loader.relation_replay_manifest:
metadata["relation_replay_manifest"] = dict(loader.relation_replay_manifest)
run_predictions = Predictions()
run_predictions.add_prediction(
dataset_name=dataset_name,
model_name=model_name,
y_pred=y_pred,
y_true=np.array([]),
partition="test",
step_idx=loader.metadata.model_step_index if loader.metadata and loader.metadata.model_step_index is not None else 0,
config_path=str(bundle_path),
fold_id="all",
metadata=metadata or None,
)
logger.success(f"Predicted with bundle: {model_name}")
if all_predictions:
return {"bundle": {"prediction": y_pred}}, run_predictions
return y_pred, run_predictions
# -----------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------
def _get_substep_from_artifact(self, artifact_id: str) -> int | None:
"""Get substep_index from artifact record.
Args:
artifact_id: Artifact ID to look up.
Returns:
substep_index or None if not found.
"""
if self.artifact_loader is None:
return None
try:
record = self.artifact_loader.get_record(artifact_id)
if record is not None:
return record.substep_index
except (KeyError, AttributeError):
pass
return None