Update model handling and data preparation

Author: seofernando25

framework/data_utils.py

@@ -1,7 +1,7 @@
 """Data loading and preprocessing utilities."""
 
 from pathlib import Path
-from typing import List, Tuple
+from typing import Any, Dict, List, Tuple
 import numpy as np
 from datasets import load_from_disk
 from sklearn.model_selection import train_test_split
@@ -133,3 +133,34 @@ def create_dataloaders(
     )
 
     return train_loader, val_loader
+
+
+def prepare_dataset(val_ratio: float = 0.1) -> Dict[str, Any]:
+    """Prepare and return the CIFAR-10 dataset"""
+    ds_dict = load_cifar10_data()
+    train_images, train_labels = prepare_data(ds_dict, "train")
+    test_images, test_labels = prepare_data(ds_dict, "test")
+
+    X_train, y_train, X_val, y_val = split_train_val(
+        train_images, train_labels, val_ratio=val_ratio
+    )
+
+    def flatten(images):
+        stacked = np.stack([np.asarray(img, dtype=np.float32) for img in images])
+        return stacked.reshape(len(images), -1)
+
+    train_flat = flatten(X_train)
+    val_flat = flatten(X_val)
+    test_flat = flatten(test_images)
+
+    return {
+        "train_images": X_train,
+        "train_labels": y_train,
+        "val_images": X_val,
+        "val_labels": y_val,
+        "test_images": test_images,
+        "test_labels": test_labels,
+        "train_flat": train_flat,
+        "val_flat": val_flat,
+        "test_flat": test_flat,
+    }

hparam_search.py

@@ -8,13 +8,12 @@
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
-from framework.data_utils import (
-    load_cifar10_data,
-    prepare_data,
-    split_train_val,
-)
+from framework.data_utils import prepare_dataset
 from framework.fitness import calculate_composite_fitness
-from models.base import get_model_by_name
+from models.cnn import CNNModel
+from models.decision_tree import DecisionTreeModel
+from models.factory import get_model_by_name
+from models.knn import KNNModel
 from search import RandomSearch
 
 RANDOM_SEED = 321
@@ -34,34 +33,6 @@ def set_seeds(seed: int):
         torch.cuda.manual_seed_all(seed)
 
 
-def prepare_dataset() -> Dict[str, Any]:
-    ds_dict = load_cifar10_data()
-    train_images, train_labels = prepare_data(ds_dict, "train")
-    test_images, test_labels = prepare_data(ds_dict, "test")
-
-    X_train, y_train, X_val, y_val = split_train_val(
-        train_images, train_labels, val_ratio=0.2
-    )
-
-    def flatten(images):
-        stacked = np.stack([np.asarray(img, dtype=np.float32) for img in images])
-        return stacked.reshape(len(images), -1)
-
-    train_flat = flatten(X_train)
-    val_flat = flatten(X_val)
-    test_flat = flatten(test_images)
-
-    return {
-        "train_images": X_train,
-        "train_labels": y_train,
-        "val_images": X_val,
-        "val_labels": y_val,
-        "test_images": test_images,
-        "test_labels": test_labels,
-        "train_flat": train_flat,
-        "val_flat": val_flat,
-        "test_flat": test_flat,
-    }
 
 
 def evaluate_model(
@@ -72,10 +43,12 @@ def evaluate_model(
     model = get_model_by_name(model_key)
 
     if model_key in {"dt", "knn"}:
+        assert isinstance(model, (DecisionTreeModel, KNNModel))
         model.create_model(**params)
         model.train(data["train_flat"], data["train_labels"])
         metrics = model.evaluate(data["val_flat"], data["val_labels"])
     elif model_key == "cnn":
+        assert isinstance(model, CNNModel)
         model.create_model(**params)
         model.train(
             data["train_images"],

models/base.py

@@ -1,11 +1,9 @@
 """Abstract interface for models used in the hyperparameter tuning framework."""
 
 from abc import ABC, abstractmethod
-from typing import Dict, Any, Literal, overload
+from typing import Dict, Any
+
 
-from models.cnn import CNNModel
-from models.decision_tree import DecisionTreeModel
-from models.knn import KNNModel
 
 from .ParamSpace import ParamSpace
 
@@ -40,30 +38,3 @@ def evaluate(self, *args: Any, **kwargs: Any) -> Dict[str, float]:
     def get_param_space(self) -> Dict[str, ParamSpace]:
         """Return the searchable hyperparameter space."""
         raise NotImplementedError
-
-
-@overload
-def get_model_by_name(model_name: Literal["dt"]) -> DecisionTreeModel:
-    ...
-
-@overload
-def get_model_by_name(model_name: Literal["knn"]) -> KNNModel:
-    ...
-
-@overload
-def get_model_by_name(model_name: Literal["cnn"]) -> CNNModel:
-    ...
-
-def get_model_by_name(model_name: Literal["dt", "knn", "cnn"]) -> KNNModel | DecisionTreeModel | CNNModel:
-    models = {
-        "dt": DecisionTreeModel,
-        "knn": KNNModel,
-        "cnn": CNNModel,
-    }
-
-    if model_name not in models:
-        raise ValueError(
-            f"Unknown model: {model_name}. Available models: {list(models.keys())}"
-        )
-
-    return models[model_name]()

models/decision_tree.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict
 
 import numpy as np
 from sklearn.metrics import classification_report, f1_score, roc_auc_score
@@ -21,23 +21,23 @@ def create_model(self, **params: Any) -> None:
         self.params.update(params)
         self.estimator = DecisionTreeClassifier(**self.params)
 
-    def train(self, X_train: List[np.ndarray], y_train: np.ndarray) -> DecisionTreeClassifier:
+    def train(self, X_train: np.ndarray, y_train: np.ndarray) -> DecisionTreeClassifier:
         if self.estimator is None:
             self.create_model()
         estimator = self.estimator
         assert estimator is not None
         estimator.fit(X_train, y_train)
         return estimator
 
-    def predict(self, X: List[np.ndarray]):
+    def predict(self, X: np.ndarray):
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."
             )
         check_is_fitted(self.estimator)
         return self.estimator.predict(X)
 
-    def predict_proba(self, X: List[np.ndarray]):
+    def predict_proba(self, X: np.ndarray):
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."
@@ -49,7 +49,7 @@ def predict_proba(self, X: List[np.ndarray]):
         check_is_fitted(self.estimator)
         return self.estimator.predict_proba(X)
 
-    def evaluate(self, X_test: List[np.ndarray], y_test: np.ndarray) -> Dict[str, float]:
+    def evaluate(self, X_test: np.ndarray, y_test: np.ndarray) -> Dict[str, float]:
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."

models/factory.py

@@ -0,0 +1,37 @@
+"""Factory function for creating model instances by name."""
+
+from typing import Literal, overload
+from models.cnn import CNNModel
+from models.decision_tree import DecisionTreeModel
+from models.knn import KNNModel
+
+@overload
+def get_model_by_name(model_name: Literal["dt"]) -> DecisionTreeModel:
+    ...
+
+
+@overload
+def get_model_by_name(model_name: Literal["knn"]) -> KNNModel:
+    ...
+
+
+@overload
+def get_model_by_name(model_name: Literal["cnn"]) -> CNNModel:
+    ...
+
+
+def get_model_by_name(model_name: Literal["dt", "knn", "cnn"]) -> KNNModel | DecisionTreeModel | CNNModel:
+
+    models = {
+        "dt": DecisionTreeModel,
+        "knn": KNNModel,
+        "cnn": CNNModel,
+    }
+
+    if model_name not in models:
+        raise ValueError(
+            f"Unknown model: {model_name}. Available models: {list(models.keys())}"
+        )
+
+    return models[model_name]()
+

models/knn.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict
 
 import numpy as np
 from sklearn.metrics import classification_report, f1_score, roc_auc_score
@@ -20,23 +20,23 @@ def create_model(self, **params: Any) -> None:
         self.params.update(params)
         self.estimator = KNeighborsClassifier(**self.params)
 
-    def train(self, X_train: List[np.ndarray], y_train: np.ndarray) -> KNeighborsClassifier:
+    def train(self, X_train: np.ndarray, y_train: np.ndarray) -> KNeighborsClassifier:
         if self.estimator is None:
             self.create_model()
         estimator = self.estimator
         assert estimator is not None
         estimator.fit(X_train, y_train)
         return estimator
 
-    def predict(self, X: List[np.ndarray]):
+    def predict(self, X: np.ndarray):
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."
             )
         check_is_fitted(self.estimator)
         return self.estimator.predict(X)
 
-    def predict_proba(self, X: List[np.ndarray]):
+    def predict_proba(self, X: np.ndarray):
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."
@@ -48,7 +48,7 @@ def predict_proba(self, X: List[np.ndarray]):
         check_is_fitted(self.estimator)
         return self.estimator.predict_proba(X)
 
-    def evaluate(self, X_test: List[np.ndarray], y_test: np.ndarray) -> Dict[str, float]:
+    def evaluate(self, X_test: np.ndarray, y_test: np.ndarray) -> Dict[str, float]:
         if self.estimator is None:
             raise RuntimeError(
                 "Estimator has not been created. Call create_model() first."

scripts/run_experiment.py

@@ -9,20 +9,19 @@
 from pathlib import Path
 from typing import Any, Dict, Literal
 
+from models.decision_tree import DecisionTreeModel
+from models.knn import KNNModel
+
 # Add project root to path for imports
 sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
 
 import numpy as np
 import torch
 
-from framework.data_utils import (
-    load_cifar10_data,
-    prepare_data,
-    split_train_val,
-)
+from framework.data_utils import prepare_dataset
 from framework.fitness import calculate_composite_fitness
-from models.base import get_model_by_name
-from models.cnn import TrainingConfig
+from models.factory import get_model_by_name
+from models.cnn import CNNModel, TrainingConfig
 from search import RandomSearch, GeneticAlgorithm, ParticleSwarmOptimization
 from dataclasses import replace
 
@@ -40,35 +39,6 @@ def set_seeds(seed: int):
         torch.cuda.manual_seed_all(seed)
 
 
-def prepare_dataset() -> Dict[str, Any]:
-    """Prepare and return the CIFAR-10 dataset."""
-    ds_dict = load_cifar10_data()
-    train_images, train_labels = prepare_data(ds_dict, "train")
-    test_images, test_labels = prepare_data(ds_dict, "test")
-
-    X_train, y_train, X_val, y_val = split_train_val(
-        train_images, train_labels, val_ratio=0.1
-    )
-
-    def flatten(images):
-        stacked = np.stack([np.asarray(img, dtype=np.float32) for img in images])
-        return stacked.reshape(len(images), -1)
-
-    train_flat = flatten(X_train)
-    val_flat = flatten(X_val)
-    test_flat = flatten(test_images)
-
-    return {
-        "train_images": X_train,
-        "train_labels": y_train,
-        "val_images": X_val,
-        "val_labels": y_val,
-        "test_images": test_images,
-        "test_labels": test_labels,
-        "train_flat": train_flat,
-        "val_flat": val_flat,
-        "test_flat": test_flat,
-    }
 
 
 def evaluate_model(
@@ -81,10 +51,12 @@ def evaluate_model(
     model = get_model_by_name(model_key)
 
     if model_key in {"dt", "knn"}:
+        assert isinstance(model, (DecisionTreeModel, KNNModel))
         model.create_model(**params)
         model.train(data["train_flat"], data["train_labels"])
         metrics = model.evaluate(data["val_flat"], data["val_labels"])
     elif model_key == "cnn":
+        assert isinstance(model, CNNModel)
         model.create_model(**params)
         default_config = TrainingConfig()
         config = replace(