"""Data loading, column-type detection, preprocessing and splitting. Design notes ------------ * Every feature column is classified as ``numeric``, ``categorical`` or ``text``. Text columns become TF-IDF vectors, categoricals are one-hot encoded and numerics are median-imputed and standardised. * The split is train / calibration / test. The calibration part is further halved (``calib_fit`` for probability calibration, ``calib_conf`` for conformal prediction and abstention thresholds) so no guarantee re-uses the data that tuned the probabilities it relies on. """ from __future__ import annotations from dataclasses import dataclass, field from pathlib import Path import numpy as np import pandas as pd from sklearn.compose import ColumnTransformer from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.impute import SimpleImputer from sklearn.model_selection import train_test_split from sklearn.pipeline import Pipeline from sklearn.preprocessing import OneHotEncoder, StandardScaler @dataclass class ColumnTypes: """Which columns are which kind of feature.""" numeric: list[str] = field(default_factory=list) categorical: list[str] = field(default_factory=list) text: list[str] = field(default_factory=list) def all(self) -> list[str]: return self.numeric + self.categorical + self.text def to_dict(self) -> dict: return {"numeric": self.numeric, "categorical": self.categorical, "text": self.text} def detect_column_types( df: pd.DataFrame, target: str | None = None, text_cols: list[str] | None = None, max_categories: int = 50, text_min_avg_tokens: float = 5.0, ) -> ColumnTypes: """Heuristically label each column (except ``target``) as numeric/categorical/text. Rules: explicit ``text_cols`` win; numeric dtypes are numeric (unless they have very few distinct integer values, which we still treat as numeric - simple and usually fine); string columns whose values average at least ``text_min_avg_tokens`` whitespace tokens are text; other string columns are categorical if they have at most ``max_categories`` distinct values and are otherwise dropped (e.g. IDs) with no complaint. """ text_cols = list(text_cols or []) types = ColumnTypes(text=[c for c in text_cols if c in df.columns]) for col in df.columns: if col == target or col in types.text: continue s = df[col] if pd.api.types.is_bool_dtype(s): types.categorical.append(col) elif pd.api.types.is_numeric_dtype(s): types.numeric.append(col) else: s_str = s.dropna().astype(str) avg_tokens = s_str.str.split().str.len().mean() if len(s_str) else 0.0 if avg_tokens >= text_min_avg_tokens: types.text.append(col) elif s_str.nunique() <= max_categories: types.categorical.append(col) # else: high-cardinality short strings (IDs, free codes) are ignored. return types def build_preprocessor(types: ColumnTypes, max_tfidf_features: int = 5000) -> ColumnTransformer: """ColumnTransformer turning a raw DataFrame into a model matrix. Output is sparse if a text column is present, dense otherwise. """ transformers = [] if types.numeric: transformers.append( ( "num", Pipeline([("impute", SimpleImputer(strategy="median")), ("scale", StandardScaler())]), types.numeric, ) ) if types.categorical: transformers.append( ( "cat", Pipeline( [ ("impute", SimpleImputer(strategy="most_frequent")), ("onehot", OneHotEncoder(handle_unknown="ignore")), ] ), types.categorical, ) ) for col in types.text: # A single column *name* (not list) hands the vectorizer a 1-D series. transformers.append( ( f"text_{col}", TfidfVectorizer( max_features=max_tfidf_features, ngram_range=(1, 2), sublinear_tf=True, min_df=2 ), col, ) ) if not transformers: raise ValueError("No usable feature columns were found.") # sparse_threshold=0 forces dense output when there is no text; with text we # keep it sparse (TF-IDF matrices are mostly zeros). return ColumnTransformer(transformers, sparse_threshold=0.3 if types.text else 0.0) def clean_text_columns(df: pd.DataFrame, types: ColumnTypes) -> pd.DataFrame: """TF-IDF chokes on NaN; replace missing text with the empty string.""" if not types.text: return df df = df.copy() for col in types.text: df[col] = df[col].fillna("").astype(str) return df def load_csv(path: str | Path, target: str, text_cols: list[str] | None = None): """Read a CSV and return ``(X_df, y, column_types)``. Rows with a missing target are dropped. ``y`` is returned as strings so that label handling is uniform (ints, bools and strings all work). """ df = pd.read_csv(path) if target not in df.columns: raise KeyError(f"target column {target!r} not in CSV columns {list(df.columns)}") df = df[df[target].notna()].reset_index(drop=True) types = detect_column_types(df, target=target, text_cols=text_cols) X = clean_text_columns(df[types.all()], types) y = df[target].astype(str).to_numpy() return X, y, types def load_embeddings(emb_path: str | Path, labels_path: str | Path): """Load precomputed embeddings (``.npy``, shape (n, d)) and their labels. ``labels_path`` may be a ``.npy`` array or a text/CSV file with one label per line (a header line is tolerated if it is not itself a label count mismatch). Returns ``(X_df, y, column_types)`` where every embedding dimension is a numeric column, so the rest of the toolkit works unchanged. """ E = np.load(emb_path) if E.ndim != 2: raise ValueError(f"embeddings must be 2-D, got shape {E.shape}") labels_path = Path(labels_path) if labels_path.suffix == ".npy": y = np.load(labels_path, allow_pickle=False) else: lines = [ln.strip() for ln in labels_path.read_text().splitlines() if ln.strip()] if len(lines) == len(E) + 1: # drop header lines = lines[1:] y = np.array(lines) if len(y) != len(E): raise ValueError(f"{len(E)} embeddings but {len(y)} labels") cols = [f"e{i}" for i in range(E.shape[1])] X = pd.DataFrame(E, columns=cols) return X, np.asarray(y).astype(str), ColumnTypes(numeric=cols) @dataclass class Split: """Index arrays for the four parts of the data.""" train: np.ndarray calib_fit: np.ndarray calib_conf: np.ndarray test: np.ndarray def sizes(self) -> dict: return {k: int(len(getattr(self, k))) for k in ("train", "calib_fit", "calib_conf", "test")} def _safe_stratify(y: np.ndarray): """Stratify only when every class has >= 2 members (sklearn requirement).""" _, counts = np.unique(y, return_counts=True) return y if counts.min() >= 2 else None def make_split( y: np.ndarray, test_size: float = 0.2, calib_size: float = 0.2, seed: int = 0 ) -> Split: """Stratified train / calib_fit / calib_conf / test split (fractions of all rows). ``calib_size`` is split equally between ``calib_fit`` and ``calib_conf``. """ idx = np.arange(len(y)) rest, test = train_test_split( idx, test_size=test_size, random_state=seed, stratify=_safe_stratify(y) ) rel_cal = calib_size / (1.0 - test_size) train, cal = train_test_split( rest, test_size=rel_cal, random_state=seed + 1, stratify=_safe_stratify(y[rest]) ) cal_fit, cal_conf = train_test_split( cal, test_size=0.5, random_state=seed + 2, stratify=_safe_stratify(y[cal]) ) return Split(np.sort(train), np.sort(cal_fit), np.sort(cal_conf), np.sort(test))