101 lines
4.3 KiB
Python
101 lines
4.3 KiB
Python
import lightgbm as lgb
|
|
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.base import BaseEstimator
|
|
from sklearn.feature_extraction.text import CountVectorizer
|
|
from sklearn.preprocessing import LabelEncoder
|
|
|
|
from research.traditional_model import TraditionalModel
|
|
|
|
|
|
class LightGBMModel(TraditionalModel):
|
|
"""LightGBM with engineered features"""
|
|
|
|
def __init__(self, config):
|
|
super().__init__(config)
|
|
# Store vectorizers and encoders to ensure consistent feature space
|
|
self.vectorizers = {}
|
|
self.label_encoders = {}
|
|
|
|
def build_model(self) -> BaseEstimator:
|
|
params = self.config.model_params
|
|
|
|
return lgb.LGBMClassifier(
|
|
n_estimators=params.get("n_estimators", 100),
|
|
max_depth=params.get("max_depth", -1),
|
|
learning_rate=params.get("learning_rate", 0.1),
|
|
num_leaves=params.get("num_leaves", 31),
|
|
subsample=params.get("subsample", 0.8),
|
|
colsample_bytree=params.get("colsample_bytree", 0.8),
|
|
random_state=self.config.random_seed,
|
|
verbose=2,
|
|
)
|
|
|
|
def prepare_features(self, X: pd.DataFrame) -> np.ndarray:
|
|
features = []
|
|
|
|
for feature_type in self.config.features:
|
|
if feature_type.value in X.columns:
|
|
column = X[feature_type.value]
|
|
|
|
if feature_type.value in ["name_length", "word_count"]:
|
|
# Numerical features
|
|
features.append(column.fillna(0).values.reshape(-1, 1))
|
|
elif feature_type.value in ["full_name", "native_name", "surname"]:
|
|
# Character-level features for names
|
|
feature_key = f"vectorizer_{feature_type.value}"
|
|
|
|
if feature_key not in self.vectorizers:
|
|
# First time - create and fit vectorizer
|
|
self.vectorizers[feature_key] = CountVectorizer(
|
|
analyzer="char", ngram_range=(2, 3), max_features=50
|
|
)
|
|
char_features = (
|
|
self.vectorizers[feature_key]
|
|
.fit_transform(column.fillna("").astype(str))
|
|
.toarray()
|
|
)
|
|
else:
|
|
# Subsequent times - use existing vectorizer
|
|
char_features = (
|
|
self.vectorizers[feature_key]
|
|
.transform(column.fillna("").astype(str))
|
|
.toarray()
|
|
)
|
|
|
|
features.append(char_features)
|
|
else:
|
|
# Categorical features
|
|
feature_key = f"encoder_{feature_type.value}"
|
|
|
|
if feature_key not in self.label_encoders:
|
|
# First time - create and fit encoder
|
|
self.label_encoders[feature_key] = LabelEncoder()
|
|
encoded = self.label_encoders[feature_key].fit_transform(
|
|
column.fillna("unknown").astype(str)
|
|
)
|
|
else:
|
|
# Subsequent times - use existing encoder
|
|
# Handle unseen labels by mapping them to a default value
|
|
column_clean = column.fillna("unknown").astype(str)
|
|
|
|
# Get the classes the encoder was trained on
|
|
known_classes = set(self.label_encoders[feature_key].classes_)
|
|
|
|
# Map unseen values to "unknown" if it exists, otherwise to the first class
|
|
if "unknown" in known_classes:
|
|
default_class = "unknown"
|
|
else:
|
|
default_class = self.label_encoders[feature_key].classes_[0]
|
|
|
|
# Replace unseen values with default
|
|
column_mapped = column_clean.apply(
|
|
lambda x: x if x in known_classes else default_class
|
|
)
|
|
|
|
encoded = self.label_encoders[feature_key].transform(column_mapped)
|
|
|
|
features.append(encoded.reshape(-1, 1))
|
|
|
|
return np.hstack(features) if features else np.array([]).reshape(len(X), 0)
|