File size: 1,398 Bytes
8bb21fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # src/models.py
from __future__ import annotations
from sklearn.ensemble import (
ExtraTreesRegressor,
RandomForestRegressor,
GradientBoostingRegressor,
StackingRegressor,
)
# Optional XGBoost (paper uses XGBR in base layer)
try:
from xgboost import XGBRegressor
HAS_XGB = True
except Exception:
HAS_XGB = False
def build_sm2_stacking(random_state: int = 42) -> StackingRegressor:
"""
SM2-style stacking:
Base: ETR, RFR, GBR, (XGBR if available)
Meta: GBR
"""
base_estimators = [
("etr", ExtraTreesRegressor(n_estimators=1000, random_state=random_state, n_jobs=-1)),
("rfr", RandomForestRegressor(n_estimators=1000, random_state=random_state, n_jobs=-1)),
("gbr", GradientBoostingRegressor(random_state=random_state)),
]
if HAS_XGB:
base_estimators.append(
("xgbr", XGBRegressor(
n_estimators=100,
max_depth=6,
learning_rate=0.1,
subsample=0.9,
reg_lambda=1.0,
random_state=random_state,
n_jobs=-1,
tree_method="hist",
))
)
meta = GradientBoostingRegressor(random_state=random_state)
return StackingRegressor(
estimators=base_estimators,
final_estimator=meta,
passthrough=True,
n_jobs=-1,
)
|