Source code for sequifier.config.metadata
"""Typed preprocessing metadata used during config resolution."""
from __future__ import annotations
import copy
import json
from typing import Any
from pydantic import (
AliasChoices,
BaseModel,
ConfigDict,
Field,
field_validator,
model_validator,
)
from sequifier.config.depth_layout import DepthLayoutRegistryModel
from sequifier.helpers import ModelWindowView, StoredWindowLayout
from sequifier.special_tokens import SPECIAL_TOKEN_IDS, validate_special_token_ids
from sequifier.typechecking import beartype
RESOLVED_ONLY_CONFIG_KEYS = {
"categorical_columns",
"real_columns",
"id_maps",
"cardinality_config",
"special_token_ids",
"storage_layout",
"window_view",
"n_classes",
"window_length",
"window_lengths",
"max_target_offset",
"stored_window_layout_version",
"depth_layouts",
"tensor_payload_version",
"prediction_aligned_splits",
"split_index",
}
[docs]class DatasetMetadata(BaseModel):
"""The stable subset of preprocessing metadata consumed by other commands."""
model_config = ConfigDict(extra="allow", populate_by_name=True)
@model_validator(mode="before")
@classmethod
def reject_legacy_placement(cls, values: Any) -> Any:
if isinstance(values, dict) and {
"window_strides",
"window_placement",
"split_context",
} & set(values):
raise ValueError(
"Preprocessing metadata uses obsolete placement settings; re-run preprocessing"
)
return values
depth_layouts: DepthLayoutRegistryModel = Field(
default_factory=DepthLayoutRegistryModel
)
tensor_payload_version: int = Field(default=1, ge=1, le=3)
split_paths: list[str] = Field(default_factory=list)
split_index: int | None = Field(default=None, ge=0)
column_data_types: dict[str, str] = Field(
default_factory=dict,
validation_alias=AliasChoices("column_data_types", "column_types"),
)
n_classes: dict[str, int] = Field(default_factory=dict)
id_maps: dict[str, dict[str | int, int]] = Field(default_factory=dict)
cardinality_config: dict[str, dict[str, Any]] = Field(default_factory=dict)
special_token_ids: dict[str, int] = Field(
default_factory=lambda: dict(SPECIAL_TOKEN_IDS.ids_by_label)
)
selected_columns_statistics: dict[str, dict[str, float]] = Field(
default_factory=dict
)
normalize_real_columns: bool = True
normalize_on_all_data: bool = False
prediction_aligned_splits: list[int] = Field(default_factory=list)
prediction_alignment_steps: int = Field(default=1, gt=0)
prediction_length: int | None = Field(default=None, gt=0)
target_offset: int | None = Field(default=None, ge=0)
window_length: int = Field(gt=0)
window_lengths: list[int] | None = None
max_target_offset: int = Field(default=1, ge=0)
stored_window_layout_version: int = Field(default=2, ge=2, le=2)
@model_validator(mode="after")
def validate_layout_metadata(self):
if self.depth_layouts and self.tensor_payload_version not in {2, 3}:
raise ValueError("Depth datasets require tensor_payload_version 2 or 3")
if self.window_lengths is not None and (
(self.split_paths and len(self.window_lengths) != len(self.split_paths))
or (
self.split_index is not None
and self.split_index >= len(self.window_lengths)
)
or any(width <= self.max_target_offset for width in self.window_lengths)
):
raise ValueError(
"Metadata window_lengths must contain one valid width per split"
)
if self.window_lengths is not None:
expected_width = self.width_for_split(self.split_index or 0)
if self.window_length != expected_width:
raise ValueError(
"Metadata window_length does not match selected split width"
)
if self.prediction_aligned_splits:
if self.prediction_length is None or self.target_offset is None:
raise ValueError(
"Aligned metadata requires prediction_length and target_offset"
)
if self.target_offset != self.max_target_offset:
raise ValueError(
"Aligned metadata requires target_offset == max_target_offset"
)
if len(self.prediction_aligned_splits) != len(
set(self.prediction_aligned_splits)
) or any(
isinstance(i, bool)
or i < 0
or (self.split_paths and i >= len(self.split_paths))
or (self.window_lengths is not None and i >= len(self.window_lengths))
for i in self.prediction_aligned_splits
):
raise ValueError("Aligned metadata has invalid split indices")
if any(
self.prediction_alignment_steps * self.prediction_length
> self.width_for_split(i) - self.max_target_offset
for i in self.prediction_aligned_splits
):
raise ValueError(
"Aligned metadata window width cannot hold all prediction steps"
)
if (
self.split_index is not None
and self.split_paths
and self.split_index >= len(self.split_paths)
):
raise ValueError("Metadata split_index is outside split_paths")
return self
def width_for_split(self, index: int) -> int:
return (
self.window_lengths[index]
if self.window_lengths is not None
else self.window_length
)
def layout_for_split(self, index: int) -> StoredWindowLayout:
return StoredWindowLayout(
self.width_for_split(index),
self.max_target_offset,
self.stored_window_layout_version,
)
@field_validator("special_token_ids")
@classmethod
@beartype
def validate_token_ids(cls, value: dict[str, int]) -> dict[str, int]:
return validate_special_token_ids(value, source="dataset metadata")
@property
@beartype
def storage_layout(self) -> StoredWindowLayout:
return StoredWindowLayout(
window_length=self.window_length,
max_target_offset=self.max_target_offset,
version=self.stored_window_layout_version,
)
[docs]@beartype
def load_dataset_metadata(path: str) -> DatasetMetadata:
"""Load and validate one preprocessing metadata JSON file."""
with open(path, "r") as file:
values: Any = json.load(file)
if not isinstance(values, dict):
raise ValueError(f"Metadata config '{path}' must contain a JSON object.")
return DatasetMetadata.model_validate(values)
[docs]@beartype
def extract_inline_metadata(
values: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any] | None]:
"""Split the historical ``skip_metadata`` representation into two mappings."""
authored = copy.deepcopy(values)
window_view = authored.get("window_view")
if "context_length" not in authored and isinstance(window_view, dict):
authored["context_length"] = window_view.get("context_length")
elif "context_length" not in authored and isinstance(window_view, ModelWindowView):
authored["context_length"] = window_view.context_length
if "target_offset" not in authored and isinstance(window_view, dict):
authored["target_offset"] = window_view.get("target_offset", 1)
elif "target_offset" not in authored and isinstance(window_view, ModelWindowView):
authored["target_offset"] = window_view.target_offset
metadata_values: dict[str, Any] = {
"split_paths": [
path
for path in (
authored.get("data_path"),
authored.get("validation_data_path"),
)
if path is not None
],
"depth_layouts": authored.get("depth_layouts", {}),
"tensor_payload_version": authored.get("tensor_payload_version", 1),
"column_data_types": authored.get("column_data_types", {}),
"n_classes": authored.get("n_classes", {}),
"id_maps": authored.get("id_maps", {}),
"cardinality_config": authored.get("cardinality_config", {}),
"special_token_ids": authored.get(
"special_token_ids", SPECIAL_TOKEN_IDS.ids_by_label
),
"selected_columns_statistics": authored.get("selected_columns_statistics", {}),
"normalize_real_columns": authored.get("normalize_real_columns", True),
"normalize_on_all_data": authored.get("normalize_on_all_data", False),
"prediction_aligned_splits": authored.get("prediction_aligned_splits", []),
"prediction_alignment_steps": authored.get("prediction_alignment_steps", 1),
"window_lengths": authored.get("window_lengths"),
"split_index": authored.get("split_index"),
"prediction_length": authored.get("prediction_length"),
"target_offset": authored.get("target_offset"),
}
storage_layout = authored.get("storage_layout")
if isinstance(storage_layout, StoredWindowLayout):
metadata_values.update(
{
"window_length": storage_layout.window_length,
"max_target_offset": storage_layout.max_target_offset,
"stored_window_layout_version": storage_layout.version,
}
)
elif isinstance(storage_layout, dict):
metadata_values.update(
{
"window_length": storage_layout.get("window_length"),
"max_target_offset": storage_layout.get("max_target_offset", 1),
"stored_window_layout_version": storage_layout.get("version", 2),
}
)
else:
metadata_values.update(
{
"window_length": authored.get("window_length"),
"max_target_offset": authored.get("max_target_offset", 1),
"stored_window_layout_version": authored.get(
"stored_window_layout_version", 2
),
}
)
for key in RESOLVED_ONLY_CONFIG_KEYS:
authored.pop(key, None)
for key in (
"selected_columns_statistics",
"normalize_real_columns",
"normalize_on_all_data",
"prediction_aligned_splits",
):
authored.pop(key, None)
if metadata_values["window_length"] is None:
return authored, None
return authored, metadata_values