wavetrainer 0.1.7__tar.gz → 0.1.8__tar.gz

{wavetrainer-0.1.7/wavetrainer.egg-info → wavetrainer-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wavetrainer
-Version: 0.1.7
+Version: 0.1.8
 Summary: A library for automatically finding the optimal model within feature and hyperparameter space.
 Home-page: https://github.com/8W9aG/wavetrainer
 Author: Will Sackfield

{wavetrainer-0.1.7 → wavetrainer-0.1.8}/setup.py

@@ -23,7 +23,7 @@ def install_requires() -> typing.List[str]:
 
 setup(
     name='wavetrainer',
-    version='0.1.7',
+    version='0.1.8',
     description='A library for automatically finding the optimal model within feature and hyperparameter space.',
     long_description=long_description,
     long_description_content_type='text/markdown',

{wavetrainer-0.1.7 → wavetrainer-0.1.8}/wavetrainer/__init__.py

@@ -2,5 +2,5 @@
 
 from .create import create
 
-__VERSION__ = "0.1.7"
+__VERSION__ = "0.1.8"
 __all__ = ("create",)

{wavetrainer-0.1.7 → wavetrainer-0.1.8}/wavetrainer/create.py

@@ -16,6 +16,7 @@ def create(
     max_train_timeout: datetime.timedelta | None = None,
     cutoff_dt: datetime.datetime | None = None,
     embedding_cols: list[list[str]] | None = None,
+    allowed_models: set[str] | None = None,
 ) -> Trainer:
     """Create a trainer."""
     return Trainer(
@@ -27,4 +28,5 @@ def create(
         max_train_timeout=max_train_timeout,
         cutoff_dt=cutoff_dt,
         embedding_cols=embedding_cols,
+        allowed_models=allowed_models,
     )

{wavetrainer-0.1.7 → wavetrainer-0.1.8}/wavetrainer/model/model_router.py

@@ -33,10 +33,13 @@ class ModelRouter(Model):
     _model: Model | None
     _false_positive_reduction_steps: int | None
 
-    def __init__(self) -> None:
+    def __init__(self, allowed_models: set[str] | None) -> None:
         super().__init__()
         self._model = None
         self._false_positive_reduction_steps = None
+        self._allowed_models = (
+            allowed_models if allowed_models is not None else set(_MODELS.keys())
+        )
 
     @classmethod
     def name(cls) -> str:
@@ -91,7 +94,12 @@ class ModelRouter(Model):
             _FALSE_POSITIVE_REDUCTION_STEPS_KEY, 0, 5
         )
         model_name = trial.suggest_categorical(
-            "model", [k for k, v in _MODELS.items() if v.supports_x(df)]
+            "model",
+            [
+                k
+                for k, v in _MODELS.items()
+                if v.supports_x(df) and k in self._allowed_models
+            ],
         )
         print(f"Using {model_name} model")
         model = _MODELS[model_name]()

{wavetrainer-0.1.7 → wavetrainer-0.1.8}/wavetrainer/trainer.py

@@ -64,6 +64,7 @@ class Trainer(Fit):
         max_train_timeout: datetime.timedelta | None = None,
         cutoff_dt: datetime.datetime | None = None,
         embedding_cols: list[list[str]] | None = None,
+        allowed_models: set[str] | None = None,
     ):
         tqdm.tqdm.pandas()
 
@@ -155,6 +156,7 @@ class Trainer(Fit):
         self._max_train_timeout = max_train_timeout
         self._cutoff_dt = cutoff_dt
         self.embedding_cols = embedding_cols
+        self._allowed_models = allowed_models
 
     def _provide_study(self, column: str) -> optuna.Study:
         storage_name = f"sqlite:///{self._folder}/{column}/{_STUDYDB_FILENAME}"
@@ -265,7 +267,7 @@ class Trainer(Fit):
                     print(f"Row weights took {time.time() - start_row_weights}")
 
                     # Create model
-                    model = ModelRouter()
+                    model = ModelRouter(self._allowed_models)
                     model.set_options(trial, x)
 
                     # Train
@@ -517,7 +519,7 @@ class Trainer(Fit):
                 reducer = CombinedReducer(self.embedding_cols)
                 reducer.load(folder)
 
-                model = ModelRouter()
+                model = ModelRouter(None)
                 model.load(folder)
 
                 selector = Selector(model)
@@ -570,7 +572,7 @@ class Trainer(Fit):
                 if not os.path.isdir(date_path):
                     continue
                 try:
-                    model = ModelRouter()
+                    model = ModelRouter(None)
                     model.load(date_path)
                     feature_importances[date_str] = model.feature_importances
                 except FileNotFoundError as exc:

{wavetrainer-0.1.7 → wavetrainer-0.1.8/wavetrainer.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wavetrainer
-Version: 0.1.7
+Version: 0.1.8
 Summary: A library for automatically finding the optimal model within feature and hyperparameter space.
 Home-page: https://github.com/8W9aG/wavetrainer
 Author: Will Sackfield