h2ofm 0.42.0__py3-none-any.whl

h2ofm/__init__.py

@@ -0,0 +1,10 @@
+# Copyright 2026 H2O.ai, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from importlib.metadata import version
+
+from . import timeh2o
+from .timeh2o import TimeH2O
+
+__version__ = version("h2ofm")
+__all__ = ["TimeH2O", "timeh2o", "__version__"]

h2ofm/timeh2o/__init__.py

@@ -0,0 +1,6 @@
+# Copyright 2026 H2O.ai, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from ._api import Forecast, TimeH2O
+
+__all__ = ["Forecast", "TimeH2O"]

h2ofm/timeh2o/_api.py

@@ -0,0 +1,630 @@
+# Copyright 2026 H2O.ai, Inc.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass
+from numbers import Integral, Real
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from pandas.api import types as ptypes
+
+from ._checkpoints import load_checkpoint, locate_checkpoint
+from ._data.preprocess import from_arrays
+from ._model.config import TimeH2OConfig
+from ._pipeline.df_utils import (
+    ITEM,
+    ROW_ORDER_FREQ,
+    STEP,
+    TIME,
+    datetimes_to_integer_times,
+    infer_freq_from_df,
+    is_integer_time_column,
+    prepare_frames,
+)
+from ._pipeline.pipeline import TimeH2OPipeline
+from ._pipeline.postprocess import clip_frame, clip_targets, nonnegative_targets
+from ._pipeline.tta import blend_items, reflect_arrays
+from ._pipeline.utils import extrapolate_quantiles, validate_quantile_levels
+from ._warnings import warn
+
+Array = np.ndarray | torch.Tensor
+DEFAULT_QUANTILES = tuple(i / 10 for i in range(1, 10))
+
+
+@dataclass(frozen=True)
+class Forecast:
+    """The forecast of one input item, as CPU NumPy arrays.
+
+    A 1D history gives ``point`` shaped (horizon,) and ``quantiles`` shaped
+    (horizon, levels); a 2D (targets, history) item keeps its target axis,
+    (targets, horizon) and (targets, horizon, levels). ``point`` holds medians,
+    not means. The last quantile axis follows ``quantile_levels``, in the
+    order they were requested.
+    """
+
+    point: np.ndarray
+    quantiles: np.ndarray
+    quantile_levels: tuple[float, ...]
+
+    @property
+    def median(self) -> np.ndarray:
+        """The point forecast under its precise name; the same array as
+        ``point``."""
+        return self.point
+
+    def quantile(self, level: float) -> np.ndarray:
+        """The forecast at one requested quantile level, shaped like
+        ``point``. Raises KeyError for a level that was not requested."""
+        if isinstance(level, bool) or not isinstance(level, Real):
+            raise TypeError(
+                f"quantile level must be a number, e.g. quantile(0.9); got {level!r}"
+            )
+        for index, requested in enumerate(self.quantile_levels):
+            if math.isclose(requested, level, rel_tol=0.0, abs_tol=1e-9):
+                return self.quantiles[..., index]
+        raise KeyError(
+            f"quantile level {level} was not requested; "
+            f"available: {list(self.quantile_levels)}"
+        )
+
+
+def _positive_int(value, name):
+    if isinstance(value, bool) or not isinstance(value, Integral) or value < 1:
+        raise ValueError(f"{name} must be a positive integer")
+    return int(value)
+
+
+def _flag(value, name):
+    if not isinstance(value, bool):
+        raise ValueError(f"{name} must be a boolean")
+    return value
+
+
+def _real_tensor(item, what: str) -> torch.Tensor:
+    if isinstance(item, torch.Tensor):
+        if item.is_complex():
+            raise ValueError(f"{what} must contain real numeric values")
+        return item.detach().to(device="cpu", dtype=torch.float32).clone()
+    array = np.asarray(item)
+    if array.dtype.kind not in "biuf":
+        raise ValueError(f"{what} must contain real numeric values")
+    return torch.tensor(np.array(array, dtype=np.float32, copy=True))
+
+
+def _series(item) -> tuple[torch.Tensor, bool]:
+    """One series as a (rows, history) tensor, and whether it was 1D."""
+    if isinstance(item, pd.DataFrame):
+        raise TypeError("use forecast_df() for dataframes")
+    if isinstance(item, (str, bytes)):
+        raise TypeError("A series must be an array, tensor or sequence of numbers")
+    item = _real_tensor(item, "Series")
+    if item.ndim not in (1, 2) or any(size == 0 for size in item.shape):
+        raise ValueError(
+            "A series must have nonempty shape (history,) or (targets, history); "
+            "use forecast_many() for several series"
+        )
+    if torch.isinf(item).any():
+        raise ValueError(
+            "Series may contain NaN for missing observations, but not infinity"
+        )
+    return item.reshape(-1, item.shape[-1]), item.ndim == 1
+
+
+def _batch(contexts) -> tuple[list[torch.Tensor], list[bool]]:
+    """Each series of a batch: a sequence, or an array iterated along axis 0."""
+    if isinstance(contexts, pd.DataFrame):
+        raise TypeError("use forecast_df() for dataframes")
+    if isinstance(contexts, (np.ndarray, torch.Tensor)):
+        if contexts.ndim not in (2, 3):
+            raise ValueError(
+                "An array of contexts must be 2D (series, history) or 3D "
+                "(series, targets, history); use forecast() for one series"
+            )
+        items = list(contexts)
+    elif isinstance(contexts, Sequence) and not isinstance(contexts, (str, bytes)):
+        if contexts and all(isinstance(v, Real) for v in contexts):
+            raise TypeError(
+                "contexts is a single series of numbers; use forecast() for one "
+                "series, or wrap it in a list"
+            )
+        items = list(contexts)
+    else:
+        raise TypeError("contexts must be a sequence of series")
+    if not items:
+        raise ValueError("contexts must contain at least one series")
+    series = [_series(item) for item in items]
+    return [t for t, _ in series], [u for _, u in series]
+
+
+_ROW_TYPES = (
+    np.ndarray,
+    torch.Tensor,
+    pd.Series,
+    pd.Index,
+    pd.Categorical,
+    list,
+    tuple,
+)
+
+
+def _covariate_row(values, name: str) -> torch.Tensor | np.ndarray:
+    """Convert non-boolean numeric rows to float32; keep other dtypes as labels."""
+    if isinstance(values, torch.Tensor):
+        values = values.detach().cpu()
+        if values.is_complex():
+            raise ValueError(f"{name} must contain real numbers or labels")
+        values = (values if values.dtype == torch.bool else values.float()).numpy()
+    if isinstance(values, np.ndarray):
+        if values.dtype.kind == "c":
+            raise ValueError(f"{name} must contain real numbers or labels")
+        if values.dtype.kind not in "iuf":
+            return values.astype(object)
+        row = torch.from_numpy(np.array(values, dtype=np.float32))
+    else:
+        series = values if isinstance(values, pd.Series) else pd.Series(values)
+        if ptypes.is_complex_dtype(series):
+            raise ValueError(f"{name} must contain real numbers or labels")
+        if not ptypes.is_numeric_dtype(series) or ptypes.is_bool_dtype(series):
+            return series.to_numpy(dtype=object)
+        row = torch.from_numpy(
+            series.to_numpy(dtype=np.float32, na_value=np.nan, copy=True)
+        )
+    if torch.isinf(row).any():
+        raise ValueError(f"{name} may contain NaN, but not infinity")
+    return row
+
+
+def _covariate_items(covariates, items, name, extra) -> list[list | None]:
+    """Parse covariate rows per item, spanning history plus extra steps."""
+    span = "history + horizon" if extra else "history"
+    if covariates is None:
+        return [None] * len(items)
+    if isinstance(covariates, (str, bytes)) or not isinstance(
+        covariates, (Sequence, np.ndarray, torch.Tensor)
+    ):
+        raise TypeError(f"{name} must be a sequence with one entry per item")
+    if len(covariates) != len(items):
+        raise ValueError(f"{name} has {len(covariates)} entries for {len(items)} items")
+    result = []
+    for index, (entry, item) in enumerate(zip(covariates, items, strict=True)):
+        if entry is None:
+            result.append(None)
+            continue
+        length = item.shape[-1] + extra
+        if isinstance(entry, (np.ndarray, torch.Tensor)) and entry.ndim == 2:
+            parts, shape = list(entry), tuple(entry.shape)
+        elif isinstance(entry, (np.ndarray, torch.Tensor)):
+            parts, shape = ([entry] if entry.ndim == 1 else None), tuple(entry.shape)
+        elif isinstance(entry, (pd.Series, pd.Index, pd.Categorical)):
+            parts, shape = [entry], (len(entry),)
+        elif isinstance(entry, Sequence) and not isinstance(entry, (str, bytes)):
+            rows = bool(entry) and all(isinstance(r, _ROW_TYPES) for r in entry)
+            parts = list(entry) if rows else [entry]
+            shape = (len(parts), len(parts[0])) if rows else (len(entry),)
+        else:
+            raise TypeError(f"{name}[{index}] must be an array or a sequence of rows")
+        if (
+            parts is None
+            or shape[-1:] != (length,)
+            or any(len(row) != length for row in parts)
+        ):
+            lengths = [len(row) for row in parts or ()]
+            got = f"rows of lengths {lengths}" if len(set(lengths)) > 1 else shape
+            raise ValueError(
+                f"{name}[{index}] must have shape ({span},) or "
+                f"(covariates, {span}) with {span} = {length}; got {got}"
+            )
+        rows = [_covariate_row(row, name) for row in parts]
+        result.append(rows or None)
+    return result
+
+
+def _warn_past_covariates_ignored(what: str) -> None:
+    warn(
+        f"This checkpoint does not use past covariates; {what} are validated "
+        "and ignored. Pass covariates known through the horizon as future "
+        "covariates instead."
+    )
+
+
+class TimeH2O:
+    """A loaded TimeH2O checkpoint.
+
+    Construct with :meth:`from_pretrained` and reuse the instance; nothing is
+    cached globally. :meth:`forecast` takes one series, :meth:`forecast_many`
+    several and :meth:`forecast_df` a long-format dataframe.
+    """
+
+    def __init__(self, pipeline: TimeH2OPipeline, batch_size: int = 256):
+        self._pipeline = pipeline
+        self._batch_size = _positive_int(batch_size, "batch_size")
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: str | Path,
+        *,
+        device: str | torch.device = "auto",
+        revision: str | None = None,
+        token: str | bool | None = None,
+        cache_dir: str | Path | None = None,
+        local_files_only: bool = False,
+        batch_size: int = 256,
+    ) -> TimeH2O:
+        """Load a checkpoint: a directory with config.json and model.safetensors.
+
+        ``pretrained_model_name_or_path`` is that directory, the
+        model.safetensors file inside it, or a Hugging Face model ID such as
+        ``"h2oai/timeh2o-v1-tiny"``, as in transformers. A path that
+        exists locally takes precedence over a model ID. From the Hub only config.json and
+        model.safetensors are fetched, into the Hugging Face cache;
+        ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` apply
+        to that download only. config.json defines the architecture; every
+        weight must match it. Pickles and remote code are never loaded.
+        'auto' selects CUDA when available, otherwise CPU. All inference uses
+        float32 weights. ``batch_size`` caps the rows (targets and covariates
+        of whole series; targets only without cross-variate attention) the
+        model runs at once; it changes speed and memory, not forecasts.
+        """
+        path = locate_checkpoint(
+            pretrained_model_name_or_path,
+            revision=revision,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+        return cls(TimeH2OPipeline(load_checkpoint(path, device)), batch_size)
+
+    @property
+    def config(self) -> TimeH2OConfig:
+        """The checkpoint's validated config.json."""
+        return self._pipeline.model.config
+
+    @property
+    def device(self) -> torch.device:
+        return self._pipeline.model.device
+
+    @property
+    def batch_size(self) -> int:
+        """The most rows run at once; set in :meth:`from_pretrained`."""
+        return self._batch_size
+
+    def _options(self, horizon, quantile_levels, tta, clip, context, ar):
+        horizon = _positive_int(horizon, "horizon")
+        flip = _flag(tta, "flip_equivariant") and not self.config.flip_equivariant
+        _flag(clip, "clip_nonnegative")
+        if ar is not None and not isinstance(ar, bool):
+            raise ValueError("autoregressive must be a boolean or None")
+        levels = [float(q) for q in quantile_levels]
+        validate_quantile_levels(levels)
+        if len(set(levels)) != len(levels):
+            raise ValueError("quantile_levels must not contain duplicates")
+        options = dict(
+            batch_size=self._batch_size,
+            context_length=(
+                None if context is None else _positive_int(context, "context_length")
+            ),
+            autoregressive=ar,
+        )
+        return horizon, levels, flip, options
+
+    @torch.no_grad()
+    def forecast(
+        self,
+        context: Array | Sequence[float],
+        horizon: int,
+        *,
+        past_covariates: Array | Sequence[float] | None = None,
+        future_covariates: Array | Sequence[float] | None = None,
+        quantile_levels: Sequence[float] = DEFAULT_QUANTILES,
+        flip_equivariant: bool = False,
+        clip_nonnegative: bool = False,
+        context_length: int | None = None,
+        autoregressive: bool | None = None,
+    ) -> Forecast:
+        """Forecast one series ``horizon`` steps ahead.
+
+        ``context`` is a 1D (history,) series or a 2D (targets, history)
+        multi-target series. ``past_covariates`` and
+        ``future_covariates`` are that series' covariates, shaped as described
+        in :meth:`forecast_many`. Returns a :class:`Forecast`; every option
+        behaves as in :meth:`forecast_many`.
+        """
+        target, univariate = _series(context)
+        return self._forecast(
+            [target],
+            [univariate],
+            [past_covariates],
+            [future_covariates],
+            horizon,
+            quantile_levels,
+            flip_equivariant,
+            clip_nonnegative,
+            context_length,
+            autoregressive,
+        )[0]
+
+    @torch.no_grad()
+    def forecast_many(
+        self,
+        contexts: Sequence[Array | Sequence[float]] | Array,
+        horizon: int,
+        *,
+        past_covariates: Sequence[Array | Sequence[float] | None] | None = None,
+        future_covariates: Sequence[Array | Sequence[float] | None] | None = None,
+        quantile_levels: Sequence[float] = DEFAULT_QUANTILES,
+        flip_equivariant: bool = False,
+        clip_nonnegative: bool = False,
+        context_length: int | None = None,
+        autoregressive: bool | None = None,
+    ) -> list[Forecast]:
+        """Forecast several series ``horizon`` steps ahead.
+
+        ``contexts`` is a sequence of series, each a 1D (history,) series or a
+        2D (targets, history) multi-target series;
+        histories may differ in length. A 2D array is a batch of 1D series and
+        a 3D array a batch of 2D series. NaN marks a missing observation; an
+        all-missing target is a zero-valued cold start. Returns one
+        :class:`Forecast` per series, in input order.
+
+        ``past_covariates`` and ``future_covariates`` hold one entry per
+        series: None, or its covariates, with NaN for a missing value. Past
+        covariates are observed over the history only, shaped (history,) or
+        (covariates, history). The timeh2o-v1 checkpoints do not use them:
+        they are validated, then ignored. Future covariates are
+        observed over the history and known through the horizon, shaped
+        (history + horizon,) or (covariates, history + horizon). An entry
+        may also be a sequence of such rows. A row with a numeric,
+        non-boolean dtype is numeric; any other row (strings, objects,
+        booleans, categoricals) holds labels, encoded as :meth:`forecast_df`
+        encodes a categorical column, with None or NaN as a missing label.
+
+        ``flip_equivariant=True``: flipping a history upside down (covariates
+        included) flips its forecast. Free for checkpoints that are
+        flip-equivariant by design (``config.flip_equivariant``); others
+        average with a flipped forecast, which roughly doubles inference time.
+        ``clip_nonnegative=True`` clamps forecasts at zero for targets with a
+        nonnegative history; it runs last, after the flip blend and quantile
+        extrapolation.
+
+        ``quantile_levels`` outside the trained 0.1..0.9 grid use
+        Gaussian-tail extrapolation. ``context_length`` may shorten the
+        checkpoint's context.
+        ``autoregressive=True`` decodes horizons beyond the checkpoint's
+        max_horizon in blocks, False in a single pass; None follows the
+        checkpoint's saved default.
+        """
+        targets, univariate = _batch(contexts)
+        size = len(targets)
+        return self._forecast(
+            targets,
+            univariate,
+            [None] * size if past_covariates is None else past_covariates,
+            [None] * size if future_covariates is None else future_covariates,
+            horizon,
+            quantile_levels,
+            flip_equivariant,
+            clip_nonnegative,
+            context_length,
+            autoregressive,
+        )
+
+    def _forecast(
+        self,
+        targets,
+        univariate,
+        past_covariates,
+        future_covariates,
+        horizon,
+        quantile_levels,
+        flip_equivariant,
+        clip_nonnegative,
+        context_length,
+        autoregressive,
+    ) -> list[Forecast]:
+        """Forecast parsed series; the covariates hold one entry per series."""
+        horizon, levels, flip, options = self._options(
+            horizon,
+            quantile_levels,
+            flip_equivariant,
+            clip_nonnegative,
+            context_length,
+            autoregressive,
+        )
+        past = _covariate_items(past_covariates, targets, "past_covariates", 0)
+        if any(entry is not None for entry in past):
+            _warn_past_covariates_ignored("past_covariates")
+        covariates = _covariate_items(
+            future_covariates, targets, "future_covariates", horizon
+        )
+        predictions = self._predict_arrays(targets, covariates, horizon, options)
+        if flip:
+            reflected = [
+                reflect_arrays(target, known)
+                for target, known in zip(targets, covariates, strict=True)
+            ]
+            predictions = blend_items(
+                predictions,
+                self._predict_arrays(
+                    [target for target, _, _ in reflected],
+                    [known for _, known, _ in reflected],
+                    horizon,
+                    options,
+                ),
+                [means for _, _, means in reflected],
+            )
+        native = self._pipeline.quantiles
+        # One call for all items: quantile maps act row by row.
+        extrapolated = extrapolate_quantiles(
+            levels, native, torch.cat(predictions).transpose(1, 2)
+        ).numpy()
+        results, start = [], 0
+        for prediction, target, squeeze in zip(
+            predictions, targets, univariate, strict=True
+        ):
+            point = prediction[:, native.index(0.5)].numpy()
+            stop = start + prediction.shape[0]
+            quantiles = extrapolated[start:stop].copy()
+            start = stop
+            if clip_nonnegative:
+                clamp = nonnegative_targets(target)
+                point = clip_targets(point, clamp)
+                quantiles = clip_targets(quantiles, clamp)
+            if squeeze:
+                point, quantiles = point[0], quantiles[0]
+            results.append(Forecast(point, quantiles, tuple(levels)))
+        return results
+
+    def _predict_arrays(self, targets, covariates, horizon, options):
+        """Return (targets, trained quantiles, horizon) forecasts per series."""
+        items = from_arrays(targets, covariates, horizon)
+        flat = self._pipeline.predict(
+            [item for parts in items for item in parts], horizon, **options
+        )
+        joined, cursor = [], 0
+        for parts in items:
+            joined.append(torch.cat(flat[cursor : cursor + len(parts)]))
+            cursor += len(parts)
+        return joined
+
+    @torch.no_grad()
+    def forecast_df(
+        self,
+        context: pd.DataFrame,
+        horizon: int | pd.DataFrame,
+        *,
+        target: str | list[str] = "target",
+        id_column: str | None = None,
+        time_column: str | None = None,
+        freq: str | int | None = None,
+        quantile_levels: Sequence[float] = DEFAULT_QUANTILES,
+        flip_equivariant: bool = False,
+        clip_nonnegative: bool = False,
+        context_length: int | None = None,
+        autoregressive: bool | None = None,
+    ) -> pd.DataFrame:
+        """Forecast from historical rows, ``horizon`` steps or rows ahead.
+
+        ``context`` contains only history, including any missing observations.
+        ``horizon`` is the number of steps to forecast per item, or a frame of
+        horizon rows. In a horizon frame, target columns must be absent or
+        entirely missing; every covariate column it has is future-known, even
+        when all its horizon values are missing, while covariates absent from
+        it are past-only and ignored by v1. It must have the same items as
+        ``context`` and an equal number of rows per item, which sets the
+        horizon. With an integer ``horizon``, every covariate is past-only.
+
+        ``target`` names one target column or a list of target columns. Omit
+        ``id_column`` for one series, and ``time_column`` to use per-item row order.
+        The time column may contain datetimes, datetime strings, or integer
+        values representing time steps; integer times are never converted to
+        dates. Supplied future times must continue history at its inferred
+        frequency, or ``freq`` when supplied: a pandas offset alias for
+        datetimes, the integer step between time values for integer times.
+        Inputs are never mutated.
+
+        Returns ``id_column`` (when given), ``time_column`` or ``step`` (1 to
+        ``horizon``), then ``target_name``, ``predictions`` (the median), and one
+        column per requested quantile, named like ``"0.1"``. Augmentation and
+        decoding options behave as in :meth:`forecast`; categorical and boolean
+        covariates remain unchanged by the flip blend.
+        """
+        horizon_df = horizon if isinstance(horizon, pd.DataFrame) else None
+        if horizon_df is None:
+            if not isinstance(horizon, Real):
+                raise TypeError(
+                    "horizon must be a positive integer or a pandas dataframe"
+                )
+            horizon = _positive_int(horizon, "horizon")
+        targets = [target] if isinstance(target, str) else target
+        history, future, frame_horizon, items = prepare_frames(
+            context, horizon_df, id_column, time_column, targets
+        )
+        dropped = [
+            c
+            for c in context.columns
+            if c not in (id_column, time_column, *targets) and c not in history
+        ]
+        if dropped:
+            _warn_past_covariates_ignored(f"past-only columns {dropped}")
+        if horizon_df is not None:
+            horizon = frame_horizon
+        horizon, levels, flip, options = self._options(
+            horizon,
+            quantile_levels,
+            flip_equivariant,
+            clip_nonnegative,
+            context_length,
+            autoregressive,
+        )
+        integer_time = time_column is not None and is_integer_time_column(
+            context[time_column]
+        )
+        if freq is not None and time_column is None:
+            raise ValueError("freq needs a time_column")
+        if isinstance(freq, bool) or (
+            integer_time and not isinstance(freq, int | None)
+        ):
+            raise ValueError("freq must be an integer step for an integer time column")
+        if isinstance(freq, int):
+            if not integer_time:
+                raise ValueError("an integer freq needs an integer time column")
+            freq = f"{_positive_int(freq, 'freq')}s"
+        output_names = {str(q) for q in levels}
+        if any(
+            isinstance(c, str) and c in output_names for c in (id_column, time_column)
+        ):
+            raise ValueError("key columns must not use quantile output names")
+        if time_column is None:
+            freq = ROW_ORDER_FREQ
+        elif freq is None:
+            freq = infer_freq_from_df(history, items, ITEM, TIME)
+        grid = self._pipeline.quantiles if flip else levels
+        predictions = self._pipeline.predict_df(
+            history,
+            future_df=future,
+            id_column=ITEM,
+            time_column=TIME,
+            target=targets,
+            horizon=horizon,
+            quantile_levels=grid,
+            freq=freq,
+            tta=flip,
+            **options,
+        )
+        if flip:
+            values = torch.from_numpy(predictions[[str(q) for q in grid]].to_numpy())
+            requested = extrapolate_quantiles(levels, grid, values).numpy()
+            for i, level in enumerate(levels):
+                predictions[str(level)] = requested[:, i]
+        if clip_nonnegative:
+            predictions = clip_frame(predictions, history, ITEM, targets, levels)
+        if time_column is None:
+            steps = predictions.groupby([ITEM, "target_name"], sort=False).cumcount()
+            predictions[STEP] = steps + 1
+        elif integer_time:
+            predictions[TIME] = datetimes_to_integer_times(predictions[TIME])
+        if id_column is not None:
+            predictions[ITEM] = items[predictions[ITEM].to_numpy()]
+        renames = {ITEM: id_column, TIME: time_column}
+        predictions = predictions.rename(
+            columns={k: v for k, v in renames.items() if v is not None}
+        )
+        keys = [
+            *([] if id_column is None else [id_column]),
+            STEP if time_column is None else time_column,
+        ]
+        return predictions[
+            [
+                *keys,
+                "target_name",
+                "predictions",
+                *(str(q) for q in levels),
+            ]
+        ]