| import json |
| import torch |
| from torch import nn |
|
|
|
|
| class Model(nn.Module): |
| """LSTM + MLP regressor that predicts the next step's numeric features. |
| |
| Each timestep is a vector of ``numeric_features`` normalized values |
| followed by an integer ticker id, which is embedded and concatenated |
| to the numeric features before the LSTM. |
| """ |
|
|
| def __init__( |
| self, |
| num_tickers, |
| numeric_features=5, |
| embedding_dim=16, |
| output_size=5, |
| ): |
| super().__init__() |
|
|
| self.numeric_features = numeric_features |
|
|
| self.ticker_embedding = nn.Embedding( |
| num_embeddings=num_tickers + 1, |
| embedding_dim=embedding_dim, |
| ) |
|
|
| self.lstm = nn.LSTM( |
| input_size=numeric_features + embedding_dim, |
| hidden_size=128, |
| num_layers=3, |
| batch_first=True, |
| ) |
|
|
| self.mlp = nn.Sequential( |
| nn.Linear(128, 128), |
| nn.GELU(), |
| nn.Linear(128, 256), |
| nn.GELU(), |
| nn.Dropout(0.2), |
| nn.Linear(256, 128), |
| nn.GELU(), |
| nn.Linear(128, output_size), |
| ) |
|
|
| def forward(self, x): |
| """Predict the next step. |
| |
| Args: |
| x: Tensor of shape (batch, sequence, numeric_features + 1) where |
| the last channel holds the ticker id. |
| |
| Returns: |
| Tensor of shape (batch, output_size). |
| """ |
| numeric = x[..., : self.numeric_features] |
| ticker_ids = x[..., self.numeric_features].long() |
| ticker_vectors = self.ticker_embedding(ticker_ids) |
| x = torch.cat([numeric, ticker_vectors], dim=-1) |
| output, _ = self.lstm(x) |
| return self.mlp(output[:, -1, :]) |
|
|
|
|
| class Normalizer(object): |
| """Applies the per-ticker mean/std normalization stored in a config file.""" |
|
|
| def __init__(self, config_path: str | None = None): |
| """Load the normalization config. |
| |
| Args: |
| config_path: Path to the JSON config. Defaults to `./config.json`. |
| """ |
| if config_path is None: |
| config_path = "./config.json" |
|
|
| with open(config_path) as f: |
| self.config = json.load(f) |
|
|
| self.ticker_ids = self.config["ticker_ids"] |
| self.mean = torch.tensor(self.config["mean"], dtype=torch.float) |
| self.std = torch.tensor(self.config["std"], dtype=torch.float) |
|
|
| def __call__(self, x, ticker_name) -> torch.Tensor: |
| """Normalize ``x`` with the statistics of ``ticker_name``.""" |
| return self.normalize(x, ticker_name) |
|
|
| def normalize(self, x, ticker_name: str) -> torch.Tensor: |
| """Return ``(x - mean) / std`` for the given ticker. |
| |
| Args: |
| x: Array-like of shape (..., num_features). Volume is expected to |
| already be log1p-transformed, as in training. |
| ticker_name: Ticker symbol present in the config's ``ticker_ids``. |
| |
| Raises: |
| KeyError: If the ticker is not in the config. |
| """ |
| ticker_id = self.ticker_ids[ticker_name] |
| x = torch.as_tensor(x, dtype=torch.float) |
| return (x - self.mean[ticker_id]) / self.std[ticker_id] |
|
|
|
|
| def load_model( |
| weights_path: str = "stock_predictor/model.pth", |
| config_path: str | None = None, |
| device: str | torch.device | None = None, |
| ) -> Model: |
| """Build a ``Model`` from the config and load trained weights into it. |
| |
| Args: |
| weights_path: Path to the saved ``state_dict``. |
| config_path: Path to the JSON config. Defaults to ``./config.json``. |
| device: Device to load onto. Defaults to CUDA if available, else CPU. |
| |
| Returns: |
| The model on ``device`` in evaluation mode. |
| """ |
| if config_path is None: |
| config_path = "./config.json" |
|
|
| if device is None: |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| device = torch.device(device) |
|
|
| with open(config_path) as f: |
| config = json.load(f) |
|
|
| model = Model( |
| num_tickers=len(config["ticker_ids"]), |
| numeric_features=len(config["mean"][0]), |
| output_size=len(config["mean"][0]), |
| ) |
| state_dict = torch.load(weights_path, map_location=device, weights_only=True) |
| model.load_state_dict(state_dict) |
| model.to(device) |
| model.eval() |
| return model |
|
|
|
|
