refactor: simplify data configuration and preprocessing steps

Author: ArturMakowski

src/mlops_with_databricks/data_preprocessing/dataclasses.py

@@ -38,27 +38,14 @@ class AdClickDataConfig:
 class ProcessedAdClickDataConfig:
     """Dataclass for the Processed Ad Click Data configuration."""
 
-    target: str = "cat__click_0"
+    target: str = "click"
     num_features: tuple[str] = ("num__age",)
     cat_features: tuple[str] = (
-        "cat__gender_Female",
-        "cat__gender_Male",
-        "cat__gender_Non-Binary",
-        "cat__device_type_Desktop",
-        "cat__device_type_Mobile",
-        "cat__device_type_Tablet",
-        "cat__ad_position_Bottom",
-        "cat__ad_position_Side",
-        "cat__ad_position_Top",
-        "cat__browsing_history_Education",
-        "cat__browsing_history_Entertainment",
-        "cat__browsing_history_News",
-        "cat__browsing_history_Shopping",
-        "cat__browsing_history_Social_Media",
-        "cat__time_of_day_Afternoon",
-        "cat__time_of_day_Evening",
-        "cat__time_of_day_Morning",
-        "cat__time_of_day_Night",
+        "cat__gender",
+        "cat__device_type",
+        "cat__ad_position",
+        "cat__browsing_history",
+        "cat__time_of_day",
     )
 
 
@@ -72,9 +59,9 @@ class DatabricksConfig:
 
 
 class LightGBMConfig(TypedDict):
-    learning_rate: str = 0.001
-    n_estimators: str = 200
-    max_depth: str = 10
+    learning_rate: float
+    n_estimators: int
+    max_depth: int
 
 
 light_gbm_config = LightGBMConfig(learning_rate=0.001, n_estimators=200, max_depth=10)

src/mlops_with_databricks/data_preprocessing/preprocess.py

@@ -10,7 +10,6 @@
 from sklearn.impute import SimpleImputer
 from sklearn.model_selection import train_test_split
 from sklearn.pipeline import Pipeline
-from sklearn.preprocessing import OneHotEncoder
 
 from mlops_with_databricks.data_preprocessing.dataclasses import AdClickDataColumns, AdClickDataConfig, DatabricksConfig
 
@@ -32,18 +31,11 @@ def load_data(self, filepath: str | Path) -> None:
     def from_pandas(cls, pandas_df: pd.DataFrame) -> "DataProcessor":
         """Create a DataProcessor object from a pandas DataFrame."""
         instance = cls()
-        instance.X = None
-        instance.y = None
-        instance.preprocessor = None
         instance.df = pandas_df
         return instance
 
     def preprocess_data(self) -> None:
-        """Preprocess the data. Fill missing values, cast types, and split features and target.
-
-        Returns:
-            tuple[pd.DataFrame, pd.Series]: Preprocessed features and target.
-        """
+        """Preprocess the data. Fill missing values, cast types, and split features and target."""
         self.df = self.df.drop(columns=[AdClickDataColumns.id, AdClickDataColumns.full_name])
         self.df = self.fill_missing_values(self.df)
         self.df[AdClickDataColumns.browsing_history] = self.df[AdClickDataColumns.browsing_history].str.replace(
@@ -56,17 +48,19 @@ def preprocess_data(self) -> None:
         categorical_transformer = Pipeline(
             steps=[
                 ("imputer", SimpleImputer(strategy="most_frequent")),
-                ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
+                # ("onehot", OneHotEncoder(handle_unknown="ignore", sparse_output=False)),
             ]
         )
 
         self.preprocessor = ColumnTransformer(
             transformers=[
                 ("num", numeric_transformer, list(AdClickDataConfig.num_features)),
-                ("cat", categorical_transformer, list(AdClickDataConfig.cat_features) + [AdClickDataConfig.target]),
+                ("cat", categorical_transformer, list(AdClickDataConfig.cat_features)),
             ]
         ).set_output(transform="pandas")
-        self.df = self.preprocessor.fit_transform(self.df)
+        preprocessed_features = self.preprocessor.fit_transform(self.df)
+        preprocessed_features["click"] = self.df[AdClickDataColumns.click].astype("int64")
+        self.df = preprocessed_features
 
     @staticmethod
     def fill_missing_values(df: pd.DataFrame) -> pd.DataFrame:

src/mlops_with_databricks/training/train.py

@@ -1,11 +1,13 @@
 # Databricks notebook source
 
+import subprocess
+
 import mlflow
 from lightgbm import LGBMClassifier
 from mlflow.models import infer_signature
 from pyspark.sql import SparkSession
 from sklearn.compose import ColumnTransformer
-from sklearn.metrics import f1_score, precision_score, recall_score, roc_auc_score
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score
 from sklearn.model_selection import GridSearchCV
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import OneHotEncoder
@@ -16,6 +18,16 @@
     light_gbm_config,
 )
 
+
+def get_git_info():
+    try:
+        sha = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("ascii").strip()
+        branch = subprocess.check_output(["git", "rev-parse", "--abbrev-ref", "HEAD"]).decode("ascii").strip()
+        return {"git_sha": sha, "branch": branch}
+    except Exception:
+        return {"git_sha": "unknown", "branch": "unknown"}
+
+
 mlflow.set_tracking_uri("databricks://dbc-643c4c2b-d6c9")
 mlflow.set_registry_uri("databricks-uc://dbc-643c4c2b-d6c9")  # It must be -uc for registering models to Unity Catalog
 
@@ -50,7 +62,7 @@
 )
 
 # Create the pipeline with preprocessing and the LightGBM regressor
-pipeline = Pipeline(steps=[("classifier", LGBMClassifier(**parameters))])
+pipeline = Pipeline(steps=[("onehot", preprocessor), ("classifier", LGBMClassifier(**parameters))])
 
 # Define parameter grid for hyperparameter tuning
 param_grid = {
@@ -60,15 +72,14 @@
 }
 
 # Perform hyperparameter tuning with GridSearchCV
-grid_search = GridSearchCV(pipeline, param_grid, cv=4, scoring="roc_auc", n_jobs=-1)
+grid_search = GridSearchCV(pipeline, param_grid, cv=4, scoring="f1", n_jobs=-1)
 
 # COMMAND ----------
 mlflow.set_experiment(experiment_name="/Shared/ad-click")
-git_sha = "ffa63b430205ff7"
 
 # Start an MLflow run to track the training process
 with mlflow.start_run(
-    tags={"git_sha": f"{git_sha}", "branch": "week2"},
+    tags=get_git_info(),
 ) as run:
     run_id = run.info.run_id
 
@@ -82,24 +93,25 @@
     precision = precision_score(y_test, y_pred)
     recall = recall_score(y_test, y_pred)
     roc_auc = roc_auc_score(y_test, y_pred)
+    accuracy = accuracy_score(y_test, y_pred)
 
     # Log parameters, metrics, and the model to MLflow
     mlflow.log_param("model_type", "LightGBM with preprocessing")
     mlflow.log_params(best_params)
-    mlflow.log_metrics({"f1": f1, "precision": precision, "recall": recall, "roc_auc": roc_auc})
-    signature = infer_signature(model_input=X_train, model_output=y_pred)
+    mlflow.log_metrics({"f1": f1, "accuracy": accuracy, "precision": precision, "recall": recall, "roc_auc": roc_auc})
+    signature = infer_signature(model_input=X_test, model_output=y_pred)
 
     dataset = mlflow.data.from_spark(train_set_spark, table_name=f"{catalog_name}.{schema_name}.train_set", version="0")
     mlflow.log_input(dataset, context="training")
 
-    mlflow.sklearn.log_model(sk_model=pipeline, artifact_path="lightgbm-pipeline-model", signature=signature)
+    mlflow.sklearn.log_model(sk_model=best_pipeline, artifact_path="lightgbm-pipeline-model", signature=signature)
 
 
 # COMMAND ----------
 model_version = mlflow.register_model(
     model_uri=f"runs:/{run_id}/lightgbm-pipeline-model",
     name=f"{catalog_name}.{schema_name}.ad_click_model_basic",
-    tags={"git_sha": f"{git_sha}"},
+    tags=get_git_info(),
 )
 
 # COMMAND ----------