Gridlock / src /preprocessing.py
parvmittal07's picture
Initial Demo complete
01eb82e
Raw
History Blame Contribute Delete
4.06 kB
"""Feature-matrix assembly with strictly train-fitted transforms.
The :class:`Preprocessor` learns everything it needs (categorical vocabularies,
embedding PCA, numeric medians) from the *training* rows only, then applies the
same transform to validation / test / future data. This is what keeps the
evaluation honest - no test statistic ever touches the fitted artefacts.
Categorical columns are emitted as pandas ``category`` dtype so gradient-boosted
trees can split on them natively (unseen test categories become NaN, which the
trees route automatically). Text embeddings are reduced with PCA fitted on the
training split.
"""
from __future__ import annotations
from dataclasses import dataclass, field
import numpy as np
import pandas as pd
from . import config as C
@dataclass
class Preprocessor:
target: str
use_embeddings: bool = True
n_emb_components: int | None = None # None -> C.EMBED_PCA_COMPONENTS
numeric_cols: list[str] = field(default_factory=list)
categorical_cols: list[str] = field(default_factory=list)
_cat_categories: dict[str, np.ndarray] = field(default_factory=dict)
_numeric_median: dict[str, float] = field(default_factory=dict)
_pca = None
_emb_mean = None
feature_names_: list[str] = field(default_factory=list)
def _select_columns(self, df: pd.DataFrame):
exclude = set()
if self.target == C.TARGET_PRIORITY:
exclude.update(C.PRIORITY_EXCLUDE_FEATURES)
num = [c for c in C.NUMERIC_FEATURES if c in df.columns and c not in exclude]
cat = [c for c in C.CATEGORICAL_FEATURES if c in df.columns and c not in exclude]
return num, cat
def fit(self, df_train: pd.DataFrame, emb_train: np.ndarray | None = None):
self.numeric_cols, self.categorical_cols = self._select_columns(df_train)
for c in self.numeric_cols:
self._numeric_median[c] = float(
pd.to_numeric(df_train[c], errors="coerce").median()
)
for c in self.categorical_cols:
cats = (
df_train[c].astype("object").where(df_train[c].notna(), np.nan)
.dropna().astype(str).unique()
)
self._cat_categories[c] = np.sort(cats)
if self.use_embeddings and emb_train is not None:
from sklearn.decomposition import PCA
target_comp = self.n_emb_components or C.EMBED_PCA_COMPONENTS
n_comp = min(target_comp, emb_train.shape[1], emb_train.shape[0] - 1)
self._emb_mean = emb_train.mean(axis=0)
self._pca = PCA(n_components=n_comp, random_state=C.RANDOM_STATE)
self._pca.fit(emb_train - self._emb_mean)
self.feature_names_ = list(self.numeric_cols) + list(self.categorical_cols)
if self._pca is not None:
self.feature_names_ += [f"emb_{i}" for i in range(self._pca.n_components_)]
return self
def transform(self, df: pd.DataFrame, emb: np.ndarray | None = None) -> pd.DataFrame:
out = {}
for c in self.numeric_cols:
col = pd.to_numeric(df[c], errors="coerce")
out[c] = col.fillna(self._numeric_median[c]).astype(np.float32)
num_df = pd.DataFrame(out, index=df.index)
cat_df = pd.DataFrame(index=df.index)
for c in self.categorical_cols:
vals = df[c].astype("object").where(df[c].notna(), np.nan).astype("string")
cat_df[c] = pd.Categorical(vals, categories=self._cat_categories[c])
parts = [num_df, cat_df]
if self._pca is not None and emb is not None:
reduced = self._pca.transform(emb - self._emb_mean).astype(np.float32)
emb_df = pd.DataFrame(
reduced, columns=[f"emb_{i}" for i in range(reduced.shape[1])],
index=df.index,
)
parts.append(emb_df)
X = pd.concat(parts, axis=1)
return X[self.feature_names_]
@property
def categorical_feature_names(self) -> list[str]:
return list(self.categorical_cols)