Skip to content

StrataFormerTrainer / StrataFormerPretrainer

StrataFormerPretrainer

Self-supervised pretrainer using Masked Bar Modeling (MBM). No labels, no manual annotation — learns from raw OHLCV data alone.

from strata import StrataFormer, StrataFormerConfig, StrataFormerPretrainer

cfg     = StrataFormerConfig(max_seq_len=256, n_layers=4)
model   = StrataFormer(cfg)
trainer = StrataFormerPretrainer(model, device="cpu", verbose=True)

pretrain(candles, epochs, seq_len, batch_size, lr, mask_prob) → StrataFormer

model = trainer.pretrain(
    candles     = raw_ohlcv_candles,   # any OHLCV list, no labels needed
    epochs      = 50,
    seq_len     = 128,                 # context window per sample
    batch_size  = 32,
    lr          = 3e-4,
    mask_prob   = 0.15,                # 15% bars masked per sequence
    weight_decay= 1e-4,
)
model.save("strata_pretrained.pt")

Training objective:

mask 15% of bars in each sequence → replace with [MASK] token
model predicts original OHLCV values of masked positions
loss = MSE(predicted, original)  [only on masked positions]

Output: trained StrataFormer in eval mode, ready for fine-tuning.

StrataFormerTrainer

Supervised fine-tuning trainer. Uses STRATA state machine as teacher — no manual labels needed.

from strata import StrataFormer, StrataFormerTrainer

trainer = StrataFormerTrainer(
    model   = pretrained_model,   # StrataFormer to fine-tune
    asset   = "AAPL",
    lr      = 1e-4,               # lower than pretraining
    device  = "cpu",
    verbose = True,
)

train(dataset, cfg, epochs, batch_size, lr, weight_decay) → StrataFormer

model = trainer.train(
    dataset      = dataset,   # StrataFormerDataset (fine-tuning mode)
    epochs       = 20,
    batch_size   = 32,
    lr           = 1e-4,
    weight_decay = 1e-4,
)
model.save("aapl_former.pt")

Loss function:

loss = CE(action_logits, action_labels)           # primary
     + 0.3 × CE(regime_logits, regime_labels)     # auxiliary
     + 0.1 × confidence_penalty                   # high conf on correct

Training schedule: CosineAnnealingLR, grad clip at 1.0

StrataFormerDataset

PyTorch Dataset supporting both pretraining and fine-tuning modes, with multi-asset support and 75% overlapping windows.

from_candles(candles, seq_len, asset, asset_id, verbose, pretrain)

from strata import StrataFormerDataset

# Fine-tuning mode (labeled)
dataset = StrataFormerDataset.from_candles(
    candles  = aapl_candles,
    seq_len  = 128,
    asset    = "AAPL",     # GUARD profile for label generation
    asset_id = 0,          # integer index for multi-asset embedding
    verbose  = True,
    pretrain = False,
)
print(dataset.action_counts())
# {"LONG": 1200, "SHORT": 380, "HOLD": 1400}

# Pretraining mode (no labels)
dataset = StrataFormerDataset.from_candles(
    candles  = any_ohlcv,
    seq_len  = 256,
    pretrain = True,
)

from_multi_asset(asset_candles, seq_len, verbose, pretrain)

dataset = StrataFormerDataset.from_multi_asset(
    asset_candles = {
        "AAPL": aapl_candles,   # asset_id=0
        "BTC":  btc_candles,    # asset_id=1
        "NVDA": nvda_candles,   # asset_id=2
        "SPY":  spy_candles,    # asset_id=3
    },
    seq_len  = 128,
    pretrain = False,
)
# Asset IDs assigned alphabetically: AAPL=0, BTC=1, NVDA=2, SPY=3
print(f"Total samples: {len(dataset)}")

Methods

dataset.action_counts()   # {"LONG": N, "SHORT": N, "HOLD": N}
dataset.regime_counts()   # {"TRENDING": N, "RANGING": N, ...}
dataset.pretrain_mode     # True = no labels (MBM mode)