import json
import os
import re
from typing import Any, Generic, Optional, TypeVar, Union
import numpy as np
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from sequifier.config.composition import (
load_composed_yaml_config,
merge_config_fragments,
)
from sequifier.config.metadata import (
DatasetMetadata,
extract_inline_metadata,
load_dataset_metadata,
)
from sequifier.helpers import (
ModelWindowView,
StoredWindowLayout,
canonicalize_polars_dtype_name,
derive_target_column_types,
metadata_config_path_from_preprocessing_data_path,
normalize_path,
resolve_window_view,
try_catch_excess_keys,
)
from sequifier.objectives import (
ALLOWED_OBJECTIVE_NAMES,
OBJECTIVE_NAME_MESSAGE,
get_objective_class,
target_offset_for_objective,
)
from sequifier.typechecking import beartype
def _is_temporal_dtype_name(dtype_name: str) -> bool:
"""Return whether a metadata dtype is a supported timestamp type."""
return dtype_name.strip().startswith(("Date", "Datetime"))
def _canonicalize_inference_dtype_name(dtype_name: str) -> str:
"""Preserve timestamp metadata while canonicalizing model-compatible dtypes."""
if _is_temporal_dtype_name(dtype_name):
return dtype_name.strip()
return canonicalize_polars_dtype_name(dtype_name)
@beartype
def _comparable_value(field: str, value: Any) -> Any:
if field == "column_data_types" and isinstance(value, dict):
return {
column: _canonicalize_inference_dtype_name(dtype)
for column, dtype in value.items()
}
return value
@beartype
def _merge_loaded_execution_values(
values: dict[str, Any], loaded: dict[str, Any], source: str
) -> None:
"""Fill omitted model fields and reject authored compatibility mismatches."""
for field, loaded_value in loaded.items():
configured_value = values.get(field)
if configured_value is not None and _comparable_value(
field, configured_value
) != _comparable_value(field, loaded_value):
raise ValueError(
f"Inference config field {field!r} does not match {source}: "
f"configured {configured_value!r}, loaded {loaded_value!r}."
)
values[field] = loaded_value
@beartype
def _execution_source(
training: Any, dataset: Any
) -> tuple[dict[str, Any], DatasetMetadata]:
interface = dataset.interface
global_spec = training.global_training
layout = interface.storage_layout
return (
{
"model_interface": dataset.model_interface,
"input_columns": list(interface.input_columns),
"target_columns": list(interface.target_columns),
"column_data_types": dict(interface.column_data_types),
"target_column_types": dict(interface.target_column_types),
"training_objective": global_spec.training_objective,
"context_length": global_spec.context_length,
"target_offset": global_spec.target_offset,
"prediction_length": interface.decoder.prediction_length,
"decoder_support": interface.decoder.support,
},
DatasetMetadata(
depth_layouts=interface.depth_layouts,
tensor_payload_version=interface.tensor_payload_version,
split_paths=interface.split_paths,
column_data_types=dict(interface.column_data_types),
n_classes=dict(interface.n_classes),
id_maps=dict(interface.id_maps),
special_token_ids=dict(interface.special_token_ids),
selected_columns_statistics=dict(interface.selected_columns_statistics),
normalize_real_columns=interface.normalize_real_columns,
prediction_aligned_splits=interface.prediction_aligned_splits,
prediction_alignment_steps=interface.prediction_alignment_steps,
window_lengths=interface.window_lengths,
prediction_length=interface.prediction_length,
target_offset=interface.target_offset,
window_length=(
interface.window_lengths[0]
if interface.window_lengths
else layout.window_length
),
max_target_offset=layout.max_target_offset,
stored_window_layout_version=layout.version,
),
)
@beartype
def _assert_metadata_matches_source(
loaded_metadata: DatasetMetadata,
selected_metadata: DatasetMetadata,
source: str,
layout_is_authoritative: bool,
) -> None:
relevant = list(loaded_metadata.column_data_types)
if loaded_metadata.depth_layouts.compatibility_signature(
relevant
) != selected_metadata.depth_layouts.compatibility_signature(relevant):
raise ValueError(f"Inference depth layouts do not match {source}")
for field in (
"column_data_types",
"n_classes",
"id_maps",
"selected_columns_statistics",
):
expected = getattr(loaded_metadata, field)
found = getattr(selected_metadata, field)
found = {key: found.get(key) for key in expected}
if _comparable_value(field, found) != _comparable_value(field, expected):
raise ValueError(
f"Inference metadata field {field!r} does not match {source}."
)
for field in ("special_token_ids", "normalize_real_columns"):
expected = getattr(loaded_metadata, field)
found = getattr(selected_metadata, field)
if found != expected:
raise ValueError(
f"Inference metadata field {field!r} does not match {source}: "
f"configured {found!r}, loaded {expected!r}."
)
if layout_is_authoritative and (
selected_metadata.max_target_offset != loaded_metadata.max_target_offset
or selected_metadata.stored_window_layout_version
!= loaded_metadata.stored_window_layout_version
):
raise ValueError(
"Inference metadata storage_layout does not match "
f"{source}: configured {selected_metadata.storage_layout!r}, loaded "
f"{loaded_metadata.storage_layout!r}."
)
@beartype
def _assert_folder_metadata_matches(
data_path: str,
selected_metadata: DatasetMetadata,
input_columns: list[str],
*,
target_columns: list[str] | None = None,
allow_wider_layout: bool = False,
) -> DatasetMetadata | None:
"""Validate a preprocessed folder's contract against selected metadata."""
if not os.path.isdir(data_path):
return None
metadata_path = os.path.join(data_path, "metadata.json")
if not os.path.isfile(metadata_path):
raise ValueError(
f"Inference data folder {data_path!r} has no metadata.json; "
"the stored-window contract cannot be validated."
)
with open(metadata_path) as metadata_file:
folder_values = json.load(metadata_file)
required_semantics = {
"id_maps",
"selected_columns_statistics",
"normalize_real_columns",
}
missing_semantics = sorted(required_semantics - folder_values.keys())
if missing_semantics:
raise ValueError(
f"Inference data folder {data_path!r} has metadata.json missing "
f"semantic fields {missing_semantics}; reprocess this data to verify "
"its category and normalization meanings."
)
folder_metadata = load_dataset_metadata(metadata_path)
mismatches = []
relevant_columns = list(dict.fromkeys(input_columns + (target_columns or [])))
expected_layout = selected_metadata.storage_layout
if folder_metadata.split_index is not None:
expected_layout = selected_metadata.layout_for_split(
folder_metadata.split_index
)
folder_layout = folder_metadata.storage_layout
layout_matches = folder_layout == expected_layout or (
allow_wider_layout
and folder_layout.window_length > expected_layout.window_length
and folder_layout.max_target_offset == expected_layout.max_target_offset
and folder_layout.version == expected_layout.version
)
if not layout_matches:
mismatches.append("storage_layout")
if (
folder_metadata.prediction_aligned_splits
!= selected_metadata.prediction_aligned_splits
or folder_metadata.prediction_length != selected_metadata.prediction_length
or folder_metadata.target_offset != selected_metadata.target_offset
or folder_metadata.prediction_alignment_steps
!= selected_metadata.prediction_alignment_steps
or (
folder_metadata.split_index is not None
and selected_metadata.split_index is not None
and folder_metadata.split_index != selected_metadata.split_index
)
):
mismatches.append("prediction_alignment")
if (
folder_metadata.tensor_payload_version
!= selected_metadata.tensor_payload_version
):
mismatches.append("tensor_payload_version")
if folder_metadata.depth_layouts.compatibility_signature(
relevant_columns
) != selected_metadata.depth_layouts.compatibility_signature(relevant_columns):
mismatches.append("depth_layouts")
selected_types = {
column: canonicalize_polars_dtype_name(
selected_metadata.column_data_types[column]
)
for column in relevant_columns
if column in selected_metadata.column_data_types
}
folder_types = {
column: canonicalize_polars_dtype_name(
folder_metadata.column_data_types[column]
)
for column in relevant_columns
if column in folder_metadata.column_data_types
}
if folder_types != selected_types:
mismatches.append("column_data_types")
categorical_columns = [
column
for column in relevant_columns
if column in selected_types and "int" in selected_types[column].lower()
]
selected_classes = {
column: selected_metadata.n_classes[column]
for column in categorical_columns
if column in selected_metadata.n_classes
}
folder_classes = {
column: folder_metadata.n_classes[column]
for column in categorical_columns
if column in folder_metadata.n_classes
}
if folder_classes != selected_classes:
mismatches.append("n_classes")
for field in ("id_maps", "selected_columns_statistics"):
selected_values = getattr(selected_metadata, field)
folder_values = getattr(folder_metadata, field)
if {
column: selected_values[column]
for column in relevant_columns
if column in selected_values
} != {
column: folder_values[column]
for column in relevant_columns
if column in folder_values
}:
mismatches.append(field)
if (
folder_metadata.normalize_real_columns
!= selected_metadata.normalize_real_columns
):
mismatches.append("normalize_real_columns")
if mismatches:
raise ValueError(
f"Inference data folder {data_path!r} has metadata fields that do not "
f"match the selected metadata: {mismatches}."
)
return folder_metadata
@beartype
def load_inferer_config(
config_path: str, args_config: dict, skip_metadata: bool
) -> "ResolvedInferenceConfig":
"""Compose inference YAML, load model defaults, and resolve metadata."""
config_values = load_composed_yaml_config(config_path)
cli_values = {
key: value for key, value in args_config.items() if key != "skip_metadata"
}
authored_values = merge_config_fragments((config_values, cli_values))
authored_values, extracted_metadata_values = extract_inline_metadata(
authored_values
)
inline_metadata_values = extracted_metadata_values if skip_metadata else None
project_root = authored_values["project_root"]
selected_metadata = None
selected_metadata_path = None
artifact_metadata_source = None
loaded_metadata_sources: list[tuple[DatasetMetadata, str, bool]] = []
dataset_name = authored_values.get("dataset")
model_paths = authored_values.get("model_path")
if isinstance(model_paths, str):
model_paths = [model_paths]
training_config_path = authored_values.get("training_config_path")
route_selected = (
dataset_name is not None or authored_values.get("model_interface") is not None
)
if training_config_path is not None and (
route_selected
or (
isinstance(model_paths, list)
and ("model_path" not in cli_values or "training_config_path" in cli_values)
and any(
isinstance(path, str) and path.lower().endswith(".pt")
for path in model_paths
)
)
):
from sequifier.config.train_config import load_train_config
training = load_train_config(training_config_path, {}, False)
datasets = training.dataset_training
if dataset_name is not None:
if dataset_name not in datasets:
raise ValueError(f"Unknown inference dataset {dataset_name!r}")
dataset = datasets[dataset_name]
elif authored_values.get("model_interface") is not None:
interface_name = authored_values["model_interface"]
dataset = next(
(
candidate
for candidate in datasets.values()
if candidate.model_interface == interface_name
),
None,
)
if dataset is None:
raise ValueError(f"Unknown model interface {interface_name!r}")
elif len(datasets) == 1:
dataset = next(iter(datasets.values()))
else:
raise ValueError(
"A dataset or model_interface selection is required for the "
"training config"
)
loaded_values, training_metadata = _execution_source(training, dataset)
source = f"training config dataset {dataset.name!r}"
_merge_loaded_execution_values(authored_values, loaded_values, source)
loaded_metadata_sources.append((training_metadata, source, True))
if dataset_name is not None:
part_name = authored_values.get("part")
if part_name is None:
if len(dataset.parts) != 1:
raise ValueError(
f"Dataset {dataset_name!r} has multiple parts; select part"
)
part_name = next(iter(dataset.parts))
authored_values["part"] = part_name
if part_name not in dataset.parts:
raise ValueError(f"Unknown inference part {dataset_name}.{part_name}")
part = dataset.parts[part_name]
selected_metadata = part.metadata
selected_metadata_path = part.metadata_config_path
import torch
from sequifier.artifacts.model_config import resolved_config_from_model_config
autoregressive_decoder_targets: set[str] = set()
if isinstance(model_paths, list):
for model_path in model_paths:
if not isinstance(model_path, str):
continue
if model_path.lower().endswith(".pt"):
payload = torch.load(
normalize_path(model_path, project_root),
map_location="cpu",
weights_only=False,
)
model_config = payload.get("model_config")
elif model_path.lower().endswith(".onnx"):
import onnx
from sequifier.model.execution_schema import (
EXECUTION_SCHEMA_KEY,
MODEL_CONFIG_KEY,
)
graph = onnx.load(
normalize_path(model_path, project_root),
load_external_data=False,
)
properties = {item.key: item.value for item in graph.metadata_props}
model_config = (
json.loads(properties[MODEL_CONFIG_KEY])
if MODEL_CONFIG_KEY in properties
else None
)
if (
any(
len(value.type.tensor_type.shape.dim) == 3
for value in graph.graph.input
)
and EXECUTION_SCHEMA_KEY not in properties
):
raise ValueError(
"Depth ONNX graphs require execution-schema metadata"
)
else:
continue
if model_config is None:
continue
if (
"model_path" in cli_values
and "training_config_path" not in cli_values
and not route_selected
):
authored_values["training_config_path"] = None
training, interface_name = resolved_config_from_model_config(
model_config,
device=str(authored_values.get("device", "cpu")),
interface_name=authored_values.get("model_interface"),
)
dataset = training.dataset_training[interface_name]
decoder = dataset.interface.decoder
branches = (
decoder.branches.values() if decoder.type == "composite" else (decoder,)
)
autoregressive_decoder_targets.update(
target
for branch in branches
if branch.type == "autoregressive_transformer"
for target in branch.target_columns
)
loaded_values, artifact_metadata = _execution_source(training, dataset)
if artifact_metadata_source is None:
artifact_metadata_source = artifact_metadata
source = f"Model artifact {model_path!r}"
_merge_loaded_execution_values(authored_values, loaded_values, source)
loaded_metadata_sources.append(
(
artifact_metadata,
source,
"storage_layout" in model_config["interfaces"][interface_name],
)
)
config = try_catch_excess_keys(config_path, InferenceConfig, authored_values)
sampled_autoregressive = autoregressive_decoder_targets.intersection(
config.sample_from_distribution_columns or ()
)
if sampled_autoregressive:
raise ValueError(
"External per-column sampling is unsupported for autoregressive "
"transformer decoder targets because later distributions depend on "
"the internally generated prefix; requested "
f"{sorted(sampled_autoregressive)!r}."
)
metadata_path = _effective_metadata_config_path(config)
artifact_only_metadata = False
if skip_metadata and inline_metadata_values is not None:
metadata = DatasetMetadata.model_validate(inline_metadata_values)
elif not skip_metadata and metadata_path is not None:
metadata = load_dataset_metadata(
normalize_path(metadata_path, config.project_root)
)
elif selected_metadata is not None:
metadata = selected_metadata
elif loaded_metadata_sources:
metadata = loaded_metadata_sources[0][0]
artifact_only_metadata = metadata is artifact_metadata_source
elif skip_metadata:
raise ValueError(
"skip_metadata requires inline storage_layout and column values, a "
"selected training dataset, or a self-describing PT artifact."
)
else:
raise ValueError(
f"Inference config '{config_path}' must define metadata_config_path "
"or preprocessing_data_path unless the PT artifact contains model "
"metadata."
)
if metadata_path is None and selected_metadata_path is not None:
config.metadata_config_path = selected_metadata_path
for loaded_metadata, source, layout_is_authoritative in loaded_metadata_sources:
_assert_metadata_matches_source(
loaded_metadata,
metadata,
source,
layout_is_authoritative,
)
return resolve_inference_config(
config, metadata, allow_wider_artifact_folder=artifact_only_metadata
)
@beartype
def _effective_metadata_config_path(config: "InferenceConfig") -> Optional[str]:
if config.metadata_config_path:
return config.metadata_config_path
if config.preprocessing_data_path:
return metadata_config_path_from_preprocessing_data_path(
config.preprocessing_data_path
)
return None
[docs]@beartype
def resolve_inference_config(
config: "InferenceConfig",
metadata: DatasetMetadata,
*,
allow_wider_artifact_folder: bool = False,
) -> "ResolvedInferenceConfig":
"""Return an inference config with all metadata-derived values populated."""
if config.prediction_alignment_steps != metadata.prediction_alignment_steps:
raise ValueError(
"Inference config prediction_alignment_steps does not match dataset "
"metadata: configured "
f"{config.prediction_alignment_steps}, "
f"metadata {metadata.prediction_alignment_steps}."
)
storage_layout = metadata.storage_layout
if storage_layout.version != 2:
raise ValueError(
"Inference requires metadata stored_window_layout_version=2, "
f"got {storage_layout.version}."
)
available_column_data_types = config.column_data_types or metadata.column_data_types
input_columns = (
[
column
for column, dtype in available_column_data_types.items()
if not _is_temporal_dtype_name(dtype)
]
if config.input_columns is None
else config.input_columns
)
relevant_columns = list(dict.fromkeys(input_columns + config.target_columns))
missing_columns = [
column
for column in relevant_columns
if column not in available_column_data_types
]
if missing_columns:
raise ValueError(
"Inference metadata is missing selected columns: " f"{missing_columns}"
)
temporal_columns = [
column
for column in relevant_columns
if _is_temporal_dtype_name(available_column_data_types[column])
]
if temporal_columns:
raise ValueError(
"Timestamp columns cannot be used as inference inputs or targets: "
f"{temporal_columns}"
)
column_data_types = {
column: available_column_data_types[column] for column in relevant_columns
}
categorical_columns = [
column
for column, type_name in column_data_types.items()
if "int" in type_name.lower() and column in input_columns
]
real_columns = [
column
for column, type_name in column_data_types.items()
if "float" in type_name.lower() and column in input_columns
]
if not categorical_columns and not real_columns:
raise ValueError("No columns found in resolved inference config")
selected_layouts = metadata.depth_layouts.relevant_layouts(input_columns)
if any(metadata.depth_layouts.is_deep_column(c) for c in config.target_columns):
raise ValueError("Depth targets are not supported")
if selected_layouts:
if config.read_format != "pt":
raise ValueError("Preprocessed depth inputs require read_format: pt")
if config.training_objective == "bert" or config.autoregressive:
raise ValueError(
"Depth inputs do not support BERT or autoregressive inference"
)
target_column_types = config.target_column_types or derive_target_column_types(
config.target_columns, column_data_types
)
window_view = ModelWindowView(
context_length=config.context_length,
objective=config.training_objective,
target_offset=target_offset_for_objective(
config.training_objective, config.target_offset
),
)
if config.data_path is None and not metadata.split_paths:
raise ValueError(
"Resolved inference config needs data_path when metadata does not "
"provide split_paths."
)
data_path = (
config.data_path or metadata.split_paths[min(2, len(metadata.split_paths) - 1)]
)
normalized_data_path = normalize_path(data_path, config.project_root)
normalized_splits = [
normalize_path(path, config.project_root) for path in metadata.split_paths
]
folder_metadata = _assert_folder_metadata_matches(
normalized_data_path,
metadata,
input_columns,
target_columns=config.target_columns,
allow_wider_layout=allow_wider_artifact_folder,
)
split_index = (
normalized_splits.index(normalized_data_path)
if normalized_data_path in normalized_splits
else metadata.split_index
)
if split_index is None and folder_metadata is not None:
split_index = folder_metadata.split_index
if split_index is not None:
storage_layout = metadata.layout_for_split(split_index)
if (
allow_wider_artifact_folder
and folder_metadata is not None
and folder_metadata.window_length > storage_layout.window_length
):
storage_layout = folder_metadata.storage_layout
resolve_window_view(storage_layout, window_view)
aligned_paths = (
{normalized_splits[index] for index in metadata.prediction_aligned_splits}
if normalized_splits
else set()
)
aligned_data_path = (
normalized_data_path in aligned_paths
or split_index in metadata.prediction_aligned_splits
)
if metadata.prediction_aligned_splits and not normalized_splits:
folder_index = None
if os.path.isdir(normalized_data_path):
match = re.search(r"-split(\d+)$", os.path.basename(normalized_data_path))
if match is not None:
folder_index = int(match.group(1))
if (
split_index is not None
and folder_index is not None
and split_index != folder_index
):
raise ValueError("Metadata split_index does not match data folder")
split_index = split_index if split_index is not None else folder_index
if split_index is None:
raise ValueError(
"Prediction aligned inference requires split_paths or a split_index "
"in metadata (or a standard -splitN data folder)"
)
aligned_data_path = split_index in metadata.prediction_aligned_splits
if aligned_data_path:
required = (
window_view.context_length
+ window_view.target_offset
+ (metadata.prediction_alignment_steps - 1)
* (metadata.prediction_length or 1)
)
if storage_layout.window_length < required:
raise ValueError(
f"Aligned inference requires window width at least {required}"
)
if aligned_data_path:
if config.model_type == "embedding" or config.autoregressive:
raise ValueError(
"Prediction aligned inference requires generative, non-autoregressive output"
)
prediction_length = config.prediction_length
if prediction_length is None:
prediction_length = get_objective_class(
config.training_objective
).default_prediction_length(window_view.context_length)
if (
metadata.target_offset != window_view.target_offset
or metadata.prediction_length != prediction_length
):
raise ValueError(
"Inference prediction view does not match preprocessing alignment"
)
if config.window_stride is not None:
raise ValueError(
"Prediction aligned inference requires window_stride: null"
)
values = config.model_dump(mode="python")
values.update(
{
"metadata_config_path": _effective_metadata_config_path(config),
"data_path": normalized_data_path,
"input_columns": input_columns,
"column_data_types": column_data_types,
"categorical_columns": categorical_columns,
"real_columns": real_columns,
"target_column_types": target_column_types,
"storage_layout": storage_layout,
"window_view": window_view,
"dataset_metadata": metadata,
"aligned_data_path": aligned_data_path,
}
)
return ResolvedInferenceConfig.model_validate(values)
_DataPathT = TypeVar("_DataPathT")
_InputColumnsT = TypeVar("_InputColumnsT")
_ColumnTypesT = TypeVar("_ColumnTypesT")
class _InferenceConfigBase(
BaseModel, Generic[_DataPathT, _InputColumnsT, _ColumnTypesT]
):
"""Shared fields and validation for authored and resolved inference config."""
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
project_root: str
preprocessing_data_path: Optional[str] = None
metadata_config_path: Optional[str] = None
model_path: Union[str, list[str]]
model_type: str
training_objective: str
data_path: _DataPathT
training_config_path: Optional[str] = None
dataset: Optional[str] = None
part: Optional[str] = None
model_interface: Optional[str] = None
read_format: str = Field(default="parquet")
write_format: str = Field(default="csv")
input_columns: _InputColumnsT
target_columns: list[str]
column_data_types: _ColumnTypesT
target_column_types: _ColumnTypesT
deterministic: bool = Field(default=False)
output_probabilities: bool = Field(default=False)
decode_categories: bool = Field(default=True)
seed: int = 1010
device: str
context_length: int = Field(gt=0)
target_offset: int = Field(default=1, ge=0)
window_stride: Optional[int] = Field(default=None, gt=0)
prediction_length: Optional[int] = None
decoder_support: int = Field(default=1, gt=0)
prediction_alignment_steps: int = Field(default=1, gt=0)
inference_batch_size: int = Field(default=1, gt=0)
sample_from_distribution_columns: Optional[list[str]] = Field(default=None)
infer_with_dropout: bool = Field(default=False)
autoregressive: bool = Field(default=False)
generation_steps: Optional[int] = Field(default=None)
@model_validator(mode="after")
@beartype
def validate_route_selection(self):
for usage, value in (
("dataset", self.dataset),
("part", self.part),
("model_interface", self.model_interface),
):
if value is not None and ("." in value or not value.isidentifier()):
raise ValueError(f"{usage} must be a valid identifier without '.'")
if self.part is not None and self.dataset is None:
raise ValueError("part selection requires dataset selection")
return self
@field_validator("input_columns", mode="before")
@classmethod
@beartype
def normalize_single_input_column(cls, value):
if isinstance(value, str):
return [value]
return value
@field_validator("training_objective")
@classmethod
@beartype
def validate_authored_training_objective(cls, value: str) -> str:
if value not in ALLOWED_OBJECTIVE_NAMES:
raise ValueError(
f"Only {OBJECTIVE_NAME_MESSAGE} are allowed, found {value}"
)
return value
@field_validator("model_type")
@classmethod
@beartype
def validate_authored_model_type(cls, value: str) -> str:
if value not in {"embedding", "generative"}:
raise ValueError("model_type must be either embedding or generative")
return value
@model_validator(mode="after")
@beartype
def validate_authored_paths(self):
ModelWindowView(
context_length=self.context_length,
objective=self.training_objective,
target_offset=target_offset_for_objective(
self.training_objective, self.target_offset
),
)
return self
[docs]class InferenceConfig(
_InferenceConfigBase[Optional[str], Optional[list[str]], Optional[dict[str, str]]]
):
"""User-authored configuration for one inference run."""
data_path: Optional[str] = None
column_data_types: Optional[dict[str, str]] = None
target_column_types: Optional[dict[str, str]] = None
[docs]class ResolvedInferenceConfig(_InferenceConfigBase[str, list[str], dict[str, str]]):
"""Internal inference config after dataset metadata has been resolved."""
categorical_columns: list[str]
real_columns: list[str]
storage_layout: StoredWindowLayout
aligned_data_path: bool = False
window_view: ModelWindowView
dataset_metadata: Optional[DatasetMetadata] = Field(default=None, exclude=True)
@model_validator(mode="before")
@classmethod
@beartype
def derive_optional_config_values(cls, values):
if not isinstance(values, dict):
return values
values = dict(values)
window_view = values.get("window_view")
if "context_length" not in values and isinstance(window_view, dict):
values["context_length"] = window_view.get("context_length")
elif "context_length" not in values and isinstance(
window_view, ModelWindowView
):
values["context_length"] = window_view.context_length
if "target_offset" not in values and isinstance(window_view, dict):
values["target_offset"] = window_view.get("target_offset", 1)
elif "target_offset" not in values and isinstance(window_view, ModelWindowView):
values["target_offset"] = window_view.target_offset
preprocessing_data_path = values.get("preprocessing_data_path")
if values.get("metadata_config_path") is None and preprocessing_data_path:
values["metadata_config_path"] = (
metadata_config_path_from_preprocessing_data_path(
preprocessing_data_path
)
)
if (
values.get("target_column_types") is None
and values.get("column_data_types") is not None
):
values["target_column_types"] = derive_target_column_types(
values.get("target_columns", []),
values["column_data_types"],
)
return values
@model_validator(mode="after")
@beartype
def validate_required_paths(self):
if self.data_path is None:
raise ValueError(
"data_path must be provided or resolved from preprocessing metadata"
)
return self
@model_validator(mode="after")
@beartype
def normalize_prediction_length(self):
if self.window_view.objective != self.training_objective:
raise ValueError(
"window_view objective must match training_objective "
f"({self.window_view.objective} != {self.training_objective})."
)
objective_class = get_objective_class(self.training_objective)
if self.prediction_length is None:
self.prediction_length = objective_class.default_prediction_length(
self.window_view.context_length
)
if objective_class.forward_looking:
decoded_length = self.window_view.context_length - self.decoder_support + 1
if not 1 <= self.prediction_length <= decoded_length:
raise ValueError(
"Forward inference prediction_length must be between 1 and "
f"context_length - decoder_support + 1 ({decoded_length})"
)
objective_class.validate_prediction_length(
self.prediction_length,
self.window_view.context_length,
usage="inference",
)
if objective_class.forward_looking:
resolve_window_view(self.storage_layout, self.window_view)
return self
@field_validator("training_objective")
@classmethod
@beartype
def validate_training_objective(cls, v):
if v not in ALLOWED_OBJECTIVE_NAMES:
raise ValueError(f"Only {OBJECTIVE_NAME_MESSAGE} are allowed, found {v}")
return v
@field_validator("model_type")
@classmethod
@beartype
def validate_model_type(cls, v: str) -> str:
if v not in [
"embedding",
"generative",
]:
raise ValueError(
f"model_type must be one of 'embedding' and 'generative, {v} isn't"
)
return v
@field_validator("output_probabilities")
@classmethod
@beartype
def validate_output_probabilities(cls, v: bool, info: Any) -> bool:
if v and info.data.get("model_type") == "embedding":
raise ValueError(
"For embedding models, 'output_probabilities' must be set to false"
)
return v
@model_validator(mode="after")
@beartype
def validate_training_config_path(self):
model_paths = (
self.model_path if isinstance(self.model_path, list) else [self.model_path]
)
if not any(path.lower().endswith(".pt") for path in model_paths):
return self
if self.training_config_path is not None and not os.path.exists(
self.training_config_path
):
raise ValueError(f"{self.training_config_path} does not exist")
return self
@field_validator("generation_steps")
@classmethod
@beartype
def validate_generation_steps(cls, v: Optional[int], info: Any) -> Optional[int]:
if v is None and info.data.get("autoregressive") is True:
raise ValueError(
"If autoregressive==True, 'generation_steps' needs to be set to an integer value."
)
if v is not None and v < 1:
raise ValueError("generation_steps must by >= 1.")
if v is not None and v > 1:
if not info.data.get("autoregressive"):
raise ValueError(
f"'generation_steps' can only be larger than 1 if 'autoregressive' is true: {info.data.get('autoregressive')}"
)
if not np.all(
np.array(sorted(info.data.get("input_columns")))
== np.array(sorted(info.data.get("target_columns")))
):
raise ValueError(
"'generation_steps' can only be larger than 1 if 'input_columns' and 'target_columns' are identical"
)
return v
@field_validator("autoregressive")
@classmethod
@beartype
def validate_autoregressive(cls, v: bool, info: Any):
if v and info.data.get("model_type") == "embedding":
raise ValueError(
"Autoregressive inference is not possible for embedding models"
)
if (
v
and info.data.get("prediction_length") is not None
and info.data.get("prediction_length") > 1
):
raise ValueError(
"Autoregressive inference is not possible for models with prediction_length > 1"
)
if v and not np.all(
np.array(sorted(info.data.get("input_columns")))
== np.array(sorted(info.data.get("target_columns")))
):
raise ValueError(
"Autoregressive inference with non-identical 'input_columns' and 'target_columns' is possible but should not be performed"
)
if (
v
and info.data.get("training_objective") is not None
and not get_objective_class(
info.data.get("training_objective")
).forward_looking
):
raise ValueError(
"Autoregressive inference is not possible with BERT-style models."
)
return v
@field_validator("data_path")
@classmethod
@beartype
def validate_data_path(cls, v: Optional[str], info: Any) -> Optional[str]:
if v is None:
return v
v2 = normalize_path(v, info.data.get("project_root"))
if not os.path.exists(v2):
raise ValueError(f"{v2} does not exist")
return v
@field_validator("read_format")
@classmethod
@beartype
def validate_read_format(cls, v: str) -> str:
if v not in ["csv", "parquet", "pt"]:
raise ValueError(
"Currently only 'csv', 'parquet' and 'pt' are supported for "
"inference input"
)
return v
@field_validator("write_format")
@classmethod
@beartype
def validate_write_format(cls, v: str) -> str:
if v not in ["csv", "parquet"]:
raise ValueError(
"Currently only 'csv' and 'parquet' are supported for inference output"
)
return v
@field_validator("target_column_types")
@classmethod
@beartype
def validate_target_column_types(
cls, v: Optional[dict], info: Any
) -> Optional[dict]:
if v is None:
return v
if not all(vv in ["categorical", "real"] for vv in v.values()):
raise ValueError(
"Target column types must be either 'categorical' or 'real'"
)
if list(v.keys()) != info.data.get("target_columns", []):
raise ValueError(
"target_columns and target_column_types must contain the same keys in the same order"
)
return v
@field_validator("column_data_types")
@classmethod
@beartype
def validate_column_types(cls, v: Optional[dict], info: Any) -> Optional[dict]:
if v is None:
return v
normalized = {
column: _canonicalize_inference_dtype_name(dtype)
for column, dtype in v.items()
}
input_columns = info.data.get("input_columns") or []
missing_input_columns = [
column for column in input_columns if column not in normalized
]
if missing_input_columns:
raise ValueError(
"column_data_types must include every input column. "
f"Missing: {missing_input_columns}"
)
return normalized
@model_validator(mode="after")
@beartype
def reject_temporal_model_columns(self):
if self.column_data_types is None:
return self
selected_columns = list(self.input_columns or []) + self.target_columns
temporal_columns = [
column
for column in dict.fromkeys(selected_columns)
if column in self.column_data_types
and _is_temporal_dtype_name(self.column_data_types[column])
]
if temporal_columns:
raise ValueError(
"Timestamp columns cannot be used as inference inputs or targets: "
f"{temporal_columns}"
)
return self
@beartype
def __init__(self, **data):
super().__init__(**data)
if self.column_data_types is None:
return
column_ordered = list(self.column_data_types.keys())
columns_ordered_filtered = [
c for c in column_ordered if c in self.target_columns
]
if not (columns_ordered_filtered == self.target_columns):
raise ValueError(f"{columns_ordered_filtered} != {self.target_columns}")