Files
drift/config/types.py
T
Mark Aron Szulyovszky 8dd2d88740 chore(Linter): reformatted code with black (#211)
* chore(Linter): reformatted code with black

* Create black.yaml
2022-02-17 19:22:17 +01:00

86 lines
2.7 KiB
Python

from pydantic import BaseModel, validator
from typing import Literal, Optional
from labeling.types import EventFilter
from models.base import Model
from data_loader.types import DataCollection, DataSource
from feature_extractors.types import FeatureExtractor, ScalerTypes
from labeling.types import EventFilter, EventLabeller
from sklearn.base import BaseEstimator
from dataclasses import dataclass
# RawConfig is needed to ensure we can declare config presets here with static typing, we then convert it to Config
class RawConfig(BaseModel):
directional_models_meta: bool
dimensionality_reduction: bool
n_features_to_select: int
expanding_window_base: bool
expanding_window_meta: bool
sliding_window_size_base: int
sliding_window_size_meta: int
retrain_every: int
scaler: Literal["normalize", "minmax", "standardize"]
assets: list[str]
target_asset: str
other_assets: list[str]
exogenous_data: list[str]
load_non_target_asset: bool
own_features: list[str]
other_features: list[str]
exogenous_features: list[str]
event_filter: Literal["none", "cusum_vol", "cusum_fixed"]
labeling: Literal["two_class", "three_class_balanced", "three_class_imbalanced"]
forecasting_horizon: int
directional_models: list[str]
meta_models: list[str]
@dataclass
class Config:
directional_models_meta: bool
dimensionality_reduction: bool
n_features_to_select: int
expanding_window_base: bool
expanding_window_meta: bool
sliding_window_size_base: int
sliding_window_size_meta: int
retrain_every: int
scaler: Literal["normalize", "minmax", "standardize"]
assets: DataCollection
target_asset: DataSource
other_assets: DataCollection
exogenous_data: DataCollection
load_non_target_asset: bool
own_features: list[tuple[str, FeatureExtractor, list[int]]]
other_features: list[tuple[str, FeatureExtractor, list[int]]]
exogenous_features: list[tuple[str, FeatureExtractor, list[int]]]
event_filter: EventFilter
labeling: EventLabeller
forecasting_horizon: int
no_of_classes: Literal["two", "three-balanced", "three-imbalanced"]
mode: Literal["training", "inference"]
directional_model: Model
meta_model: Model
@validator("directional_model", "meta_model")
def check_model(cls, v):
assert isinstance(v, BaseEstimator)
return v
@validator("event_filter")
def check_event_filter(cls, v):
assert isinstance(v, EventFilter)
return v
@validator("labeling")
def check_labeling(cls, v):
assert isinstance(v, EventLabeller)
return v
# class Config:
# arbitrary_types_allowed = True