| """Online stateful inference wrapper for the causal GRU. |
| |
| Carries the GRU hidden state across steps so each prediction is O(1) work, |
| which is what makes the model fit the 1-core / 60-minute submission budget. |
| Used by both local stepwise evaluation and the packaged solution.py. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Optional |
|
|
| import numpy as np |
| import torch |
|
|
| from src.models.sequence_models import CausalGRUForecaster |
|
|
|
|
| class GRUStatefulPredictionModel: |
| """Wrap a CausalGRUForecaster as a competition PredictionModel.""" |
|
|
| def __init__(self, model: CausalGRUForecaster, n_features: int = 32): |
| self.model = model.eval() |
| self.n_features = n_features |
| self.current_seq: Optional[int] = None |
| self.h: Optional[torch.Tensor] = None |
|
|
| def reset(self): |
| self.current_seq = None |
| self.h = None |
|
|
| @torch.no_grad() |
| def predict(self, data_point): |
| if self.current_seq != data_point.seq_ix: |
| self.current_seq = data_point.seq_ix |
| self.h = None |
| state = torch.from_numpy(np.asarray(data_point.state, dtype=np.float32)) |
| pred, self.h = self.model.step(state, self.h) |
| if not data_point.need_prediction: |
| return None |
| return pred.detach().numpy().astype(np.float32) |
|
|
|
|
| def build_model_from_config(params: dict, model_type: str = "CausalGRUForecaster"): |
| if model_type == "CausalTCN": |
| from src.models.tcn import CausalTCN |
| return CausalTCN( |
| n_features=params.get("n_features", 32), |
| d_model=params.get("d_model", 192), |
| n_layers=params.get("n_layers", 6), |
| kernel_size=params.get("kernel_size", 3), |
| dropout=params.get("dropout", 0.1), |
| head_hidden=params.get("head_hidden"), |
| ) |
| return CausalGRUForecaster( |
| n_features=params.get("n_features", 32), |
| d_model=params.get("d_model", 256), |
| n_layers=params.get("n_layers", 2), |
| dropout=params.get("dropout", 0.1), |
| head_hidden=params.get("head_hidden"), |
| rnn_type=params.get("rnn_type", "gru"), |
| ) |
|
|
|
|
| def load_gru_checkpoint(path: str | Path, map_location: str = "cpu"): |
| """Reconstruct a forecaster (GRU/LSTM/TCN) from a saved checkpoint.""" |
| ckpt = torch.load(str(path), map_location=map_location, weights_only=False) |
| model_cfg = ckpt.get("config", {}).get("model", {}) |
| params = model_cfg.get("params", {}) |
| model = build_model_from_config(params, model_cfg.get("type", "CausalGRUForecaster")) |
| model.load_state_dict(ckpt["model_state_dict"]) |
| model.eval() |
| return model |
|
|
