StockAI

Highlights

• ~11,500 lines of Python across 42 modules covering ingest, feature engineering, regime detection, modeling, validation, and live serving.
• 23 scrapers across 5 kinds of source (market data, Reddit sentiment, SEC filings, macro indicators, per-ticker news), each on its own schedule.
• 148 features, with correlation-based selection dropping redundant inputs on each run.
• A MultiHeadLSTM with one shared encoder predicts 10 timeframes, from minutes to weeks.
• Pre-training runs in three phases (daily, hourly, minute). The encoder carries over between phases with new per-timeframe heads each time, so the longer timeframes give the shorter ones a head start.
• Feature attention (EMA α=0.15, clipped to [0.5, 2.0]) rescales input weights as regimes shift, without a full retrain.
• HMM regime detection limits some heads to certain regimes and drives per-regime scoring, so a model has to do well in more than one regime to go live.
• Every reported number comes from Purged K-Fold cross-validation with embargo gaps (Lopez de Prado), which keeps future data out of training.
• Retraining with rollback: a new model must beat the live one on out-of-fold direction accuracy and per-regime accuracy, and a live model that keeps getting worse is rolled back.
• Three storage tiers: SQLite for live reads, Parquet for training, and compressed archives for long backtests.
• FastAPI + WebSocket dashboard streams predictions, feature attention weights, current regime label, and validator stats live.

Stack

Problem

Approach

class FeatureAttention:
    def __init__(self, n_features: int, alpha: float = 0.15,
                 lo: float = 0.5, hi: float = 2.0):
        self.alpha, self.lo, self.hi = alpha, lo, hi
        self.ema  = np.ones(n_features)     # per-feature weight
        self.hist = np.ones(n_features)     # EMA of |contribution|

    def update(self, contrib: np.ndarray) -> np.ndarray:
        """
        contrib[i] = running importance score for feature i
        (e.g. |grad_i * x_i| averaged over the last batch).
        Returns the weight to scale feature i by on the next pass.
        """
        self.hist = (1 - self.alpha) * self.hist \
                  + self.alpha * np.abs(contrib)
        mean   = self.hist.mean() + 1e-9
        target = self.hist / mean                      # >1 = important
        self.ema = np.clip(target, self.lo, self.hi)   # bound rescale
        return self.ema

def purged_kfold_indices(n: int, k: int,
                         label_h: int, embargo: int):
    fold_size = n // k
    for i in range(k):
        test_start = i * fold_size
        test_end   = test_start + fold_size
        test_idx   = np.arange(test_start, test_end)

        # Purge: drop training points whose label window
        # (t .. t + label_h) overlaps the test fold.
        train_idx = np.arange(n)
        train_idx = train_idx[(train_idx + label_h < test_start)
                              | (train_idx >= test_end)]

        # Embargo: drop the first 'embargo' training points
        # immediately after the test fold to block post-test leakage.
        train_idx = train_idx[(train_idx < test_start)
                              | (train_idx >= test_end + embargo)]

        yield train_idx, test_idx

def promote_if_better(candidate, incumbent,
                      folds, regimes) -> bool:
    c_dir = oof_direction_accuracy(candidate, folds)
    i_dir = oof_direction_accuracy(incumbent, folds)

    c_reg = regime_stratified_accuracy(candidate, folds, regimes)
    i_reg = regime_stratified_accuracy(incumbent, folds, regimes)

    # Must beat incumbent on BOTH axes before replacing.
    if c_dir > i_dir and min(c_reg.values()) > min(i_reg.values()):
        deploy(candidate)
        log.info("PROMOTE: dir %.3f -> %.3f, "
                 "worst-regime %.3f -> %.3f",
                 i_dir, c_dir,
                 min(i_reg.values()), min(c_reg.values()))
        return True
    return False


def live_rollback_check(live, window: int = 5):
    # If live direction accuracy has fallen below the prior
    # checkpoint's validation score for 'window' consecutive
    # days, roll back to the prior checkpoint.
    recent = live.recent_daily_accuracy(window)
    if all(a < live.prior_checkpoint.val_accuracy for a in recent):
        live.rollback()

Figures

Why these choices

MultiHeadLSTM (shared encoder, 10 heads)

One encoder learns the shared sequence representation, and each timeframe's head specializes. It's cheaper to train than 10 separate models, and because the timeframes share one representation, their predictions can be combined.

Closed-loop feature attention (EMA α=0.15)

Regime shifts change which inputs matter. The attention module tracks each feature's contribution and rescales its input weight on the next pass, without a full retrain. Weights stay within [0.5, 2.0], so no input can drop to zero or take over.

HMM regime detection (hmmlearn)

Direction accuracy averaged across regimes is misleading: a model can score 55% overall while hitting 70% in one regime and 40% in another. The HMM label lets the validator score each regime separately, and it limits some heads to the regimes they were trained for.

Purged K-Fold + embargo (Lopez de Prado)

Standard K-Fold leaks future data into training when targets span overlapping bars. Purging drops the overlapping observations, and the embargo gap blocks leakage after the test window. Without both, backtest numbers come out too high.

3-tier storage (SQLite / Parquet / compressed archive)

Live reads come from SQLite (fast, single writer). Training reads columnar Parquet (10-100× faster scans on numeric features). Long backtests and audits read compressed archives.

Retrain-with-rollback

Every new model has to beat the live one on two metrics before it replaces it, and if the live model keeps getting worse, the system rolls back to the previous checkpoint.

FastAPI + WebSocket dashboard

Predictions, attention weights, the regime label, and validator stats stream live over a WebSocket, so I can check a bad hour of signals on the dashboard without opening a notebook.