src.core.models

🦉 AthenaAI is an AI-powered ensemble machine-learning trading bot for PocketOption, built on BinaryOptionsToolsV2 → Get it on GitLab
💬 Need help? Join the Chipa Discord · ⚡ Go no-code with ChipaEditor

Source: src/core/models.py. Defines EnsemblePredictor, the bot’s prediction brain: three online scikit-learn models that adapt in real time plus two batch models trained on accumulated data (and, when trained via train.py, an optional GPU MLP — see torch-model.md).

The module imports scikit-learn inside a try/except. If sklearn is missing, SKLEARN_OK = False, a warning is printed at import, and the predictor degrades to a rule-based fallback (see _fallback_predict).

class EnsemblePredictor

__init__(self)

No parameters. Builds:

• scaler: StandardScaler shared by all models.
• models — three online learners, each a dict {"name", "clf", "accuracy_ema"} starting at accuracy_ema = 0.5:
  • "SGD": SGDClassifier(loss="modified_huber", penalty="l2", alpha=1e-4, warm_start=True, random_state=42)
  • "PA": SGDClassifier(loss="hinge", penalty=None, learning_rate="pa1", eta0=1.0, warm_start=True, random_state=42) (a Passive-Aggressive setup)
  • "NB": GaussianNB()
• _batch_models — two batch learners in the same dict shape:
  • "GBM": GradientBoostingClassifier(n_estimators=200, max_depth=4, learning_rate=0.1, subsample=0.8, random_state=42)
  • "RF": RandomForestClassifier(n_estimators=200, max_depth=8, random_state=42, n_jobs=-1)
• State: _batch_fitted=False, _fitted=False, sample buffers _X_buffer/_y_buffer (online) and _batch_X/_batch_y (cumulative, never auto-cleared), _classes = np.array([0, 1]).

Label convention everywhere: 1 = CALL would win (price rose), 0 = PUT would win (price fell).

If sklearn is missing, models = [] and scaler = None; only the fallback path works.

add_sample(self, features: np.ndarray, label: int) -> None

Appends to both the online buffers and the cumulative batch buffers. No training happens here; no return value. Note the batch buffers grow unboundedly unless trimmed externally — src/bot.py trims them to the last 500 samples before periodic batch retrains.

partial_fit(self) -> None

Trains the online models on everything buffered since the last call, then clears the buffer. No-op if sklearn is unavailable or the buffer is empty.

Steps:

1. Stacks the buffer into X (np.vstack) and y.
2. Dimension-change reset: if already fitted and scaler.n_features_in_ != X.shape[1], logs Feature dimension changed (a → b) — resetting. and calls self.__init__() — all learned state (online and batch models, scaler, buffers) is discarded and the method returns. This is the recovery path when the feature-vector length changes (e.g. code updates adding features).
3. Fits the scaler (fit on first call, partial_fit thereafter) and transforms X.
4. For each online model: clf.partial_fit(X_scaled, y, classes=[0, 1]) if supported, else full clf.fit. If already fitted, updates the vote weight: accuracy_ema = 0.9 * accuracy_ema + 0.1 * batch_accuracy (accuracy measured on this same batch — an in-sample estimate).
5. Sets _fitted = True, clears the online buffers, and logs e.g. Models updated | acc EMAs: {'SGD': '0.612', ...}.

train_batch_models(self) -> None

Trains GBM and RF on all samples currently in _batch_X/_batch_y.

• Requires ≥ 100 samples; otherwise logs Not enough data for batch models (N) and returns.
• Transforms with the existing scaler (the scaler must already be fitted — call partial_fit() first, as both the bot and trainer do).
• For each batch model: fit, then set accuracy_ema to the training-set accuracy (not EMA-smoothed). A per-model exception is caught and logged as a warning; training continues with the others.
• Sets _batch_fitted = True. Logs 🧠 Training batch models (GBM + RF) on N samples … and per-model accuracy.

Persistence

save_brain(self, path="athena_po_brain.pkl") -> None

Pickles a dict to path:

{
  "scaler": StandardScaler,
  "models": [(name, clf, accuracy_ema), ...],        # online
  "batch_models": [(name, clf, accuracy_ema), ...],  # GBM, RF, optionally MLP
  "batch_fitted": bool,
  "fitted": bool,
}

If _fitted is False, logs No trained models to save. and writes nothing. Side effects: writes the file (no directory creation — the caller must ensure it exists; Trainer.run() does), logs 🧠 Brain saved to <path>. I/O errors are not caught here (the bot wraps its call in try/except).

load_brain(self, path="athena_po_brain.pkl") -> bool

Returns False if the file does not exist. Unpickles and restores the scaler, the three online models (zipped positionally against the fresh defaults), and — importantly — rebuilds _batch_models from the file entirely, so a brain trained by train.py that contains an extra "MLP" entry (TorchMLPClassifier) is loaded with all its models. Restores _batch_fitted and _fitted. Any exception is caught, logged as Failed to load brain: <e>, and False is returned. On success logs the loaded accuracy EMAs and returns True.

Security note: this is pickle.load — only load brain files you created yourself.

predict(self, features: np.ndarray) -> tuple[Direction, float]

Returns (Direction, confidence) where confidence is a fraction in roughly [0.5, 0.78].

Fallback conditions (any of these routes to _fallback_predict): sklearn missing, not yet fitted, scaler not fitted, len(features) differs from scaler.n_features_in_, or the scaler transform raises.

Vote weighting (the important internals):

1. Each online model votes its predict_proba P(CALL) (or a hard 0/1 for models without predict_proba, i.e. the hinge-loss PA model), weighted by its accuracy_ema. A per-model predict_proba exception yields a neutral 0.5 vote.
2. If _batch_fitted, each batch model (GBM, RF, and MLP when present) votes the same way but with weight accuracy_ema * 0.75 — demoted in v6 (previously 2.0) because the batch models overfit and produced anti-correlated confidence.
3. p = weighted_call / (total_weight + 1e-10).
4. Confidence clamp: p is clipped to [0.22, 0.78] — live data showed >78% confidence was overfit noise (~50% win rate).
5. If p >= 0.5 return (Direction.CALL, p), else (Direction.PUT, 1 - p) — so the returned confidence is always ≥ 0.5 and ≤ 0.78.

Consequence for tuning: setting PO_MIN_CONF above 0.78 means the bot can never trade.

_fallback_predict(features) -> tuple[Direction, float] (static)

Rule-based prediction used before training or without sklearn. Reads features[14] (RSI), features[12] (MACD histogram), features[10] (SMA cross) with safe defaults if the vector is short. Scores: RSI < 30 → +0.3, RSI > 70 → −0.3; MACD hist ± → ±0.2; SMA cross ± → ±0.15. Confidence = 0.5 + min(|score|, 0.45). Positive score → CALL, otherwise PUT (a zero score returns PUT).

Usage example

import numpy as np
from src.core.models import EnsemblePredictor

ens = EnsemblePredictor()
for x, y in zip(X_train, y_train):        # x: np.ndarray, y: 0/1
    ens.add_sample(x, int(y))
ens.partial_fit()          # online models + scaler
ens.train_batch_models()   # GBM + RF (needs >= 100 samples)

direction, conf = ens.predict(X_test[0])
ens.save_brain("athena_brain.pkl")

See also

• Ensemble concept
• torch-model.md — the optional MLP member
• Persistence guide

💬 Unsure how this interacts with the rest of the configuration? Ask in the Chipa Discord — or prototype strategies no-code with ChipaEditor.