Spaces:
Sleeping
Sleeping
File size: 3,679 Bytes
8fe2df2 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
Model construction and training.
build_pipeline_map() -> dict of {ticker: (feature_type, fresh_model_clone)}
train_models() -> dict of {ticker: (feature_type, fitted_model)}
"""
import pandas as pd
from sklearn.ensemble import (
RandomForestClassifier,
HistGradientBoostingClassifier,
VotingClassifier,
)
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.base import clone
from core.config import TICKERS, DATA_DIR, PIPELINE_MAP
from core.features import extract_semantic_features, extract_sequential_features
# ── Base model templates ─────────────────────────────────────────────────────
_RF = RandomForestClassifier(
random_state=42, n_estimators=300, max_depth=8,
min_samples_leaf=5, n_jobs=-1,
)
_GBM = HistGradientBoostingClassifier(
random_state=42, max_iter=300, l2_regularization=1.0, max_depth=8,
)
_ENSEMBLE = VotingClassifier(
estimators=[("rf", _RF), ("gbm", _GBM)], voting="soft",
)
_LR_PIPELINE = Pipeline([
("scaler", StandardScaler()),
("lr", LogisticRegression(C=0.1, max_iter=1000)),
])
_KBEST15_LR = Pipeline([
("scaler", StandardScaler()),
("kbest", SelectKBest(f_classif, k=15)),
("lr", LogisticRegression(C=0.1, max_iter=1000)),
])
_MODEL_TEMPLATES = {
"ensemble": _ENSEMBLE,
"lr_pipeline": _LR_PIPELINE,
"kbest15_lr": _KBEST15_LR,
}
# ── Public API ───────────────────────────────────────────────────────────────
def build_pipeline_map():
"""
Return a dict of {ticker: (feature_type, fresh_model_clone)}.
Uses PIPELINE_MAP from config to look up the architecture per ticker.
"""
result = {}
for ticker in TICKERS:
feat_type, model_key = PIPELINE_MAP[ticker]
result[ticker] = (feat_type, clone(_MODEL_TEMPLATES[model_key]))
return result
def train_models(log_fn=None):
"""
Load parquet data, extract features, and train all ticker models.
Parameters
----------
log_fn : callable(str), optional
Logging function (e.g. logger.info). Falls back to print.
Returns
-------
dict { ticker: (feature_type, fitted_model) }
"""
if log_fn is None:
log_fn = print
pipeline_map = build_pipeline_map()
models = {}
for ticker in TICKERS:
fpath = DATA_DIR / f"{ticker}_minute.parquet"
if not fpath.exists():
log_fn(f"[{ticker}] Parquet file not found: {fpath}")
continue
df = pd.read_parquet(fpath)
df["date"] = pd.to_datetime(df["date"])
df.set_index("date", inplace=True)
df.sort_index(inplace=True)
# Keep at most 300 trading days
unique_days = df.index.normalize().unique()
if len(unique_days) > 300:
df = df[df.index.normalize().isin(unique_days[-300:])]
feat_type, clf = pipeline_map[ticker]
if feat_type == "semantic":
X, y, _ = extract_semantic_features(df)
else:
X, y, _ = extract_sequential_features(df)
if X is None or X.empty:
log_fn(f"[{ticker}] Feature extraction returned empty!")
continue
clf.fit(X, y)
models[ticker] = (feat_type, clf)
log_fn(f"[{ticker}] Model trained ({feat_type}) | {len(X)} samples")
return models
|