Improved a few data types.

Made a few data types more specific and accurate. Improved some
descriptions of CLI arguments.

Author: Sebastijan-Dominis

docs/configs/interim_data.md

@@ -30,25 +30,25 @@ The purpose of the interim stage is to *transform raw data into memory-optimized
 
 |Field|Type|Description|
 |:----|:---|:----------|
-|`data`|object|Dataset identity and output configuration|
-|`data_schema`|object|Expected schema of the interim dataset|
-|`raw_data_version`|string|Raw dataset version used as input|
-|`cleaning`|object|Column normalization rules|
-|`invariants`|object|Column-level validation and filtering rules|
-|`drop_duplicates`|bool|Whether duplicate rows should be removed|
-|`drop_missing_ints`|bool|Whether rows with missing integer values should be dropped|
-|`min_rows`|int|Minimum number of rows required after cleaning|
-|`lineage`|object|Metadata describing config provenance|
+|`data`|object|Dataset identity and output configuration.|
+|`data_schema`|object|Expected schema of the interim dataset.|
+|`raw_data_version`|string|Raw dataset version used as input.|
+|`cleaning`|object|Column normalization rules.|
+|`invariants`|object|Column-level validation and filtering rules.|
+|`drop_duplicates`|boolean|Whether duplicate rows should be removed.|
+|`drop_missing_ints`|boolean|Whether rows with missing integer values should be dropped. Defaults to True if not specified.|
+|`min_rows`|integer|Minimum number of rows required after cleaning. Defaults to 0 if not specified.|
+|`lineage`|object|Metadata describing config provenance.|
 
 ### `data`
 
 Defines metadata for the produced interim dataset.
 
 |Field|Type|Description|
 |:----|:---|:----------|
-|`name`|string|Dataset name|
-|`version`|string|Dataset version (`v{integer}` format)|
-|`output`|object|Output storage configuration|
+|`name`|string|Dataset name.|
+|`version`|string|Dataset version (`v{integer}` format).|
+|`output`|object|Output storage configuration.|
 
 ### `data.output`
 

ml/config/schemas/model_specs.py

@@ -27,8 +27,6 @@ class TaskType(StrEnum):
 
     classification = "classification"
     regression = "regression"
-    ranking = "ranking"
-    time_series = "time_series"
 
 class TaskConfig(BaseModel):
     """Task type metadata for model execution and validation logic."""
@@ -165,14 +163,14 @@ def validate_filters_based_on_enabled(cls, v, info):
     @field_validator("include_in_model", mode="after")
     @classmethod
     def validate_include_in_model_based_on_enabled(cls, v, info):
-        """Validate model-include flag consistency with segmentation state.
+        """Validate include_in_model flag consistency with segmentation state.
 
         Args:
-            v: Include-in-model flag value.
+            v: Include_in_model flag value.
             info: Pydantic field-validation context.
 
         Returns:
-            bool: Validated include-in-model flag.
+            bool: Validated include_in_model flag.
         """
 
         enabled = info.data.get("enabled", False)
@@ -191,8 +189,26 @@ class FeatureSetConfig(BaseModel):
 
     name: str
     version: str
-    data_format: str
-    file_name: str
+    data_format: Literal['csv', 'parquet', 'json', 'arrow'] = "parquet"
+    file_name: str = "features.parquet"
+
+    @field_validator("version", mode="before")
+    @classmethod
+    def validate_version_format(cls, v):
+        """Ensure feature set version follows ``v{number}`` convention.
+
+        Args:
+            v: Feature set version value.
+
+        Returns:
+            str: Validated feature set version string.
+        """
+
+        if not isinstance(v, str) or not v.startswith("v") or not v[1:].isdigit():
+            msg = f"Version must be in format 'v{{number}}', e.g. 'v1', 'v2', etc. Got '{v}'."
+            logger.error(msg)
+            raise ConfigError(msg)
+        return v
 
 class SplitConfig(BaseModel):
     """Train/validation/test split strategy settings."""
@@ -224,6 +240,24 @@ class PipelineConfig(BaseModel):
     version: str
     path: str
 
+    @field_validator("version", mode="before")
+    @classmethod
+    def validate_version_format(cls, v):
+        """Ensure pipeline config version follows ``v{number}`` convention.
+
+        Args:
+            v: Pipeline config version value.
+
+        Returns:
+            str: Validated pipeline config version string.
+        """
+
+        if not isinstance(v, str) or not v.startswith("v") or not v[1:].isdigit():
+            msg = f"Version must be in format 'v{{number}}', e.g. 'v1', 'v2', etc. Got '{v}'."
+            logger.error(msg)
+            raise ConfigError(msg)
+        return v
+
 class ScoringConfig(BaseModel):
     """Metric-scoring policy and optional thresholds."""
 
@@ -342,7 +376,7 @@ class ExplainabilityConfig(BaseModel):
     top_k: int = 20
     methods: ExplainabilityMethodsConfig = Field(default_factory=ExplainabilityMethodsConfig)
 
-DATA_TYPE = Literal["tabular", "time-series"]
+DATA_TYPE = Literal["tabular"]
 
 class ModelSpecsLineageConfig(BaseModel):
     """Lineage metadata for model specification creation."""
@@ -385,6 +419,24 @@ class ModelSpecs(BaseModel):
 
     model_config = ConfigDict(extra="forbid", populate_by_name=True)
 
+    @field_validator("version", mode="before")
+    @classmethod
+    def validate_version_format(cls, v):
+        """Ensure model version follows ``v{number}`` convention.
+
+        Args:
+            v: Model version value.
+
+        Returns:
+            str: Validated model version string.
+        """
+
+        if not isinstance(v, str) or not v.startswith("v") or not v[1:].isdigit():
+            msg = f"Version must be in format 'v{{number}}', e.g. 'v1', 'v2', etc. Got '{v}'."
+            logger.error(msg)
+            raise ConfigError(msg)
+        return v
+
     @model_validator(mode="after")
     def validate_task_target_consistency(self):
         """Enforce consistency between task type and target class metadata.

ml/data/config/schemas/shared.py

@@ -1,20 +1,42 @@
 """Shared schema blocks reused by interim and processed data configs."""
 
+import logging
 from typing import Literal
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
+from ml.exceptions import ConfigError
+
+logger = logging.getLogger(__name__)
 
 class Output(BaseModel):
     """Output artifact settings for persisted datasets."""
 
     path_suffix: str = Field(..., description="Suffix for the output data file path, which will be combined with the data name and version to create the full path.")
-    format: Literal["parquet"] = Field("parquet", description="Format to save the interim data.")
-    compression: Literal['snappy', 'gzip', 'brotli', 'lz4', 'zstd'] | None = Field(None, description="Compression method to use when saving the data (default: 'snappy').")
+    format: Literal["parquet"] = Field("parquet", description="Format to save the interim data (default: 'parquet').")
+    compression: Literal['snappy', 'gzip', 'brotli', 'lz4', 'zstd'] | None = Field("snappy", description="Compression method to use when saving the data (default: 'snappy').")
 
 class DataInfo(BaseModel):
     """Dataset identity and output target metadata."""
 
     name: str = Field(..., description="Name of the data being processed.")
     version: str = Field(..., description="Version of the interim data being created.")
     output: Output
+
+    @field_validator("version", mode="before")
+    @classmethod
+    def validate_version_format(cls, v):
+        """Ensure dataset version follows ``v{number}`` convention.
+
+        Args:
+            v: Dataset version value.
+
+        Returns:
+            str: Validated dataset version string.
+        """
+
+        if not isinstance(v, str) or not v.startswith("v") or not v[1:].isdigit():
+            msg = f"Version must be in format 'v{{number}}', e.g. 'v1', 'v2', etc. Got '{v}'."
+            logger.error(msg)
+            raise ConfigError(msg)
+        return v

ml/feature_freezing/freeze_strategies/tabular/config/models.py

@@ -30,7 +30,7 @@ class DatasetConfig(BaseModel):
     ref: str = Field("data/processed", description="Reference path for the dataset, e.g., 'data/processed'")
     name: str = Field(..., description="Name of the dataset, e.g., 'hotel_bookings'")
     version: str = Field(..., description="Version of the dataset, e.g., 'v1'")
-    format: Literal["csv", "parquet"]
+    format: Literal["csv", "parquet", "json", "arrow"]
     merge_key: str | list[str] = Field(
         "row_id", description="Key(s) to merge datasets on, default is 'row_id'"
     )
@@ -44,6 +44,24 @@ class DatasetConfig(BaseModel):
         "data.{format}", description="Suffix for the dataset file, supports {format} placeholder"
     )
 
+    @field_validator("version", mode="before")
+    @classmethod
+    def validate_version_format(cls, v):
+        """Ensure dataset version follows ``v{number}`` convention.
+
+        Args:
+            v: Dataset version value.
+
+        Returns:
+            str: Validated dataset version string.
+        """
+
+        if not isinstance(v, str) or not v.startswith("v") or not v[1:].isdigit():
+            msg = f"Version must be in format 'v{{number}}', e.g. 'v1', 'v2', etc. Got '{v}'."
+            logger.error(msg)
+            raise ConfigError(msg)
+        return v
+
     @field_validator("merge_key", mode="before")
     def ensure_merge_key_list(cls, v):
         """Always convert merge_key to a list internally."""
@@ -101,7 +119,7 @@ class StorageConfig(BaseModel):
     """Snapshot storage format and compression settings."""
 
     format: Literal["parquet"]
-    compression: str | None = "snappy"
+    compression: Literal['snappy', 'gzip', 'brotli', 'lz4', 'zstd'] | None = Field("snappy", description="Compression method to use when saving the data (default: 'snappy').")
 
 class LineageConfig(BaseModel):
     """Lineage metadata for feature registry config provenance."""

ml/pipelines/models.py

@@ -21,7 +21,7 @@ class PipelineConfig(BaseModel):
     version: str = Field(..., description="Version of the pipeline (e.g., v1)")
     description: str | None = Field(None, description="Optional description of the pipeline")
     steps: list[str] = Field(..., description="Ordered list of pipeline steps")
-    assumptions: dict = Field(..., description="Assumptions about supported tasks and categorical handling")
+    assumptions: dict[str, bool] = Field(..., description="Assumptions about supported tasks and categorical handling")
     lineage: LineageConfig
 
     @field_validator("version")

scripts/generators/generate_fake_data.py

@@ -102,7 +102,7 @@ def parse_args():
         "--model-path",
         type=str,
         default=None,
-        help="Path to a pre-trained CTGAN model to load. If not provided, uses a default path."
+        help="Path to a pre-trained CTGAN model to load. If not provided, trains a new model."
     )
 
     parser.add_argument(