feat(ensemble): add TFT inference adapter with sequence buffering
Implements TftInferenceAdapter that wraps the Temporal Fusion Transformer for ensemble prediction. Unlike single-step models (DQN, PPO), TFT requires a sequence of observations before running inference. The adapter maintains an internal VecDeque buffer that collects sequence_length feature vectors, then constructs static/historical/future tensors for the TFT forward pass. Median quantile is mapped to directional signal via smooth saturating function; IQR provides confidence estimation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
319
ml/src/ensemble/adapters/mamba2.rs
Normal file
319
ml/src/ensemble/adapters/mamba2.rs
Normal file
@@ -0,0 +1,319 @@
|
||||
//! Mamba2 SSM inference adapter for ensemble prediction
|
||||
//!
|
||||
//! Wraps a loaded Mamba2 state-space model with a sequence buffer.
|
||||
//! Feature vectors are zero-padded to `d_model` dimensions and accumulated
|
||||
//! in a ring buffer. Once `sequence_length` frames are collected the SSM
|
||||
//! runs a forward pass and the last timestep's sigmoid output is mapped to
|
||||
//! a directional signal + confidence for ensemble aggregation.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
|
||||
use crate::ensemble::inference_adapter::{
|
||||
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
|
||||
};
|
||||
use crate::mamba::{Mamba2Config, Mamba2SSM};
|
||||
use crate::{MLError, MLResult};
|
||||
|
||||
/// Inference adapter that wraps a Mamba2 SSM for ensemble prediction.
|
||||
///
|
||||
/// Maintains an internal sequence buffer of zero-padded feature vectors.
|
||||
/// When the buffer reaches `sequence_length`, each call to [`predict`]
|
||||
/// builds a `[1, seq_len, d_model]` tensor, runs the SSM forward pass,
|
||||
/// and converts the sigmoid output of the last timestep into:
|
||||
///
|
||||
/// - **direction**: `2 * prob - 1` mapping `[0,1] -> [-1,1]`
|
||||
/// - **confidence**: `|prob - 0.5| * 2` (highest at extremes, lowest at 0.5)
|
||||
#[allow(missing_debug_implementations)]
|
||||
pub struct Mamba2InferenceAdapter {
|
||||
model: Mutex<Mamba2SSM>,
|
||||
buffer: Mutex<VecDeque<Vec<f64>>>,
|
||||
sequence_length: usize,
|
||||
d_model: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
// SAFETY: Mamba2SSM internally uses candle tensors which are Send+Sync.
|
||||
// The Mutex provides exclusive access for inference calls and buffer mutation.
|
||||
unsafe impl Send for Mamba2InferenceAdapter {}
|
||||
unsafe impl Sync for Mamba2InferenceAdapter {}
|
||||
|
||||
impl Mamba2InferenceAdapter {
|
||||
/// Create a new Mamba2 inference adapter from configuration.
|
||||
///
|
||||
/// Initializes a fresh Mamba2 SSM with random weights on the best
|
||||
/// available device (CUDA GPU if available, otherwise CPU).
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `config` - Mamba2 model configuration (d_model determines padding width)
|
||||
/// * `sequence_length` - Number of feature vectors to buffer before running inference
|
||||
pub fn new(config: Mamba2Config, sequence_length: usize) -> MLResult<Self> {
|
||||
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
||||
let d_model = config.d_model;
|
||||
let model = Mamba2SSM::new(config, &device)?;
|
||||
Ok(Self {
|
||||
model: Mutex::new(model),
|
||||
buffer: Mutex::new(VecDeque::with_capacity(sequence_length)),
|
||||
sequence_length,
|
||||
d_model,
|
||||
device,
|
||||
})
|
||||
}
|
||||
|
||||
/// Zero-pad a feature vector to `d_model` dimensions.
|
||||
///
|
||||
/// Copies up to `d_model` values from the input. If the input is shorter
|
||||
/// than `d_model`, the remaining elements are zero-padded.
|
||||
fn pad_to_d_model(&self, values: &[f64]) -> Vec<f64> {
|
||||
let mut padded = vec![0.0f64; self.d_model];
|
||||
let copy_len = values.len().min(self.d_model);
|
||||
for i in 0..copy_len {
|
||||
if let Some(v) = values.get(i) {
|
||||
if let Some(slot) = padded.get_mut(i) {
|
||||
*slot = *v;
|
||||
}
|
||||
}
|
||||
}
|
||||
padded
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelInferenceAdapter for Mamba2InferenceAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"MAMBA-2"
|
||||
}
|
||||
|
||||
fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Zero-pad the feature vector to d_model dimensions
|
||||
let padded = self.pad_to_d_model(&features.values);
|
||||
|
||||
// Push to buffer (acquire and release buffer lock before model lock)
|
||||
let buffer_ready = {
|
||||
let mut buf = self
|
||||
.buffer
|
||||
.lock()
|
||||
.map_err(|e| MLError::LockError(format!("Mamba2 buffer lock poisoned: {e}")))?;
|
||||
buf.push_back(padded);
|
||||
// Trim to sequence_length if overflowed
|
||||
while buf.len() > self.sequence_length {
|
||||
buf.pop_front();
|
||||
}
|
||||
buf.len() >= self.sequence_length
|
||||
};
|
||||
// buffer lock is dropped here
|
||||
|
||||
if !buffer_ready {
|
||||
// Not enough data yet — return neutral prediction
|
||||
let latency_us = start.elapsed().as_micros() as u64;
|
||||
return Ok(EnsemblePrediction {
|
||||
model_name: "MAMBA-2".to_string(),
|
||||
direction: 0.0,
|
||||
confidence: 0.0,
|
||||
metadata: PredictionMeta {
|
||||
latency_us,
|
||||
quantiles: None,
|
||||
attention_weights: None,
|
||||
q_values: None,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
// Build input tensor from buffer snapshot
|
||||
let flat_data = {
|
||||
let buf = self
|
||||
.buffer
|
||||
.lock()
|
||||
.map_err(|e| MLError::LockError(format!("Mamba2 buffer lock poisoned: {e}")))?;
|
||||
let mut data = Vec::with_capacity(self.sequence_length * self.d_model);
|
||||
for frame in buf.iter() {
|
||||
data.extend_from_slice(frame);
|
||||
}
|
||||
data
|
||||
};
|
||||
// buffer lock is dropped here
|
||||
|
||||
// Tensor shape: [1, seq_len, d_model]
|
||||
let input = Tensor::from_vec(
|
||||
flat_data,
|
||||
(1, self.sequence_length, self.d_model),
|
||||
&self.device,
|
||||
)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to create Mamba2 input tensor: {e}")))?;
|
||||
|
||||
// Cast to F64 (Mamba2 uses DType::F64 internally)
|
||||
let input = input
|
||||
.to_dtype(DType::F64)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to cast input to F64: {e}")))?;
|
||||
|
||||
// Run forward pass (needs &mut self)
|
||||
let mut model = self
|
||||
.model
|
||||
.lock()
|
||||
.map_err(|e| MLError::LockError(format!("Mamba2 model lock poisoned: {e}")))?;
|
||||
let output = model.forward(&input)?;
|
||||
|
||||
// Output shape: [1, seq_len, 1]
|
||||
// Squeeze batch dim → [seq_len, 1], then squeeze trailing dim → [seq_len]
|
||||
let squeezed = output
|
||||
.squeeze(0)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze batch dim: {e}")))?;
|
||||
let squeezed = squeezed
|
||||
.squeeze(1)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze output dim: {e}")))?;
|
||||
|
||||
// Extract last timestep prediction
|
||||
let all_values: Vec<f64> = squeezed
|
||||
.to_vec1()
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to extract Mamba2 output: {e}")))?;
|
||||
|
||||
let last_idx = all_values
|
||||
.len()
|
||||
.checked_sub(1)
|
||||
.ok_or_else(|| MLError::InferenceError("Mamba2 produced empty output".to_string()))?;
|
||||
let prob = all_values
|
||||
.get(last_idx)
|
||||
.copied()
|
||||
.ok_or_else(|| MLError::InferenceError("Last timestep index out of bounds".to_string()))?;
|
||||
|
||||
// Map sigmoid output [0, 1] to direction [-1, 1] and confidence [0, 1]
|
||||
let direction = (2.0 * prob - 1.0).clamp(-1.0, 1.0);
|
||||
let confidence = ((prob - 0.5).abs() * 2.0).clamp(0.0, 1.0);
|
||||
|
||||
let latency_us = start.elapsed().as_micros() as u64;
|
||||
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "MAMBA-2".to_string(),
|
||||
direction,
|
||||
confidence,
|
||||
metadata: PredictionMeta {
|
||||
latency_us,
|
||||
quantiles: None,
|
||||
attention_weights: None,
|
||||
q_values: None,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
let buf_ok = self
|
||||
.buffer
|
||||
.lock()
|
||||
.map(|buf| buf.len() >= self.sequence_length)
|
||||
.unwrap_or(false);
|
||||
let model_ok = self.model.lock().is_ok();
|
||||
buf_ok && model_ok
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn test_config() -> Mamba2Config {
|
||||
Mamba2Config {
|
||||
d_model: 32,
|
||||
d_state: 8,
|
||||
d_head: 8,
|
||||
num_heads: 2,
|
||||
expand: 2,
|
||||
num_layers: 1,
|
||||
max_seq_len: 8,
|
||||
dropout: 0.0,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
const TEST_SEQ_LEN: usize = 4;
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_adapter_creation() {
|
||||
let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN)
|
||||
.expect("Failed to create Mamba2 adapter");
|
||||
assert_eq!(adapter.model_name(), "MAMBA-2");
|
||||
// Buffer is empty at creation → not ready
|
||||
assert!(!adapter.is_ready());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_adapter_buffers_and_predicts() {
|
||||
let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN)
|
||||
.expect("Failed to create Mamba2 adapter");
|
||||
|
||||
// Feed sequence_length feature vectors
|
||||
for i in 0..TEST_SEQ_LEN {
|
||||
let fv = FeatureVector {
|
||||
values: vec![0.1 * (i as f64 + 1.0); 51],
|
||||
timestamp: 1700000000_000_000 + i as i64,
|
||||
};
|
||||
let pred = adapter.predict(&fv).expect("predict failed");
|
||||
|
||||
if i < TEST_SEQ_LEN - 1 {
|
||||
// Not enough data yet — neutral prediction
|
||||
assert_eq!(
|
||||
pred.direction, 0.0,
|
||||
"Expected neutral direction before buffer full"
|
||||
);
|
||||
assert_eq!(
|
||||
pred.confidence, 0.0,
|
||||
"Expected zero confidence before buffer full"
|
||||
);
|
||||
} else {
|
||||
// Buffer is full — should produce a real prediction
|
||||
assert!(
|
||||
pred.direction >= -1.0 && pred.direction <= 1.0,
|
||||
"direction {} out of [-1,1]",
|
||||
pred.direction
|
||||
);
|
||||
assert!(
|
||||
pred.confidence >= 0.0 && pred.confidence <= 1.0,
|
||||
"confidence {} out of [0,1]",
|
||||
pred.confidence
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// Adapter should now be ready
|
||||
assert!(adapter.is_ready());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_adapter_deterministic() {
|
||||
let adapter = Mamba2InferenceAdapter::new(test_config(), TEST_SEQ_LEN)
|
||||
.expect("Failed to create Mamba2 adapter");
|
||||
|
||||
// Fill the buffer with identical feature vectors
|
||||
for _ in 0..TEST_SEQ_LEN {
|
||||
let fv = FeatureVector {
|
||||
values: vec![0.1; 51],
|
||||
timestamp: 1700000000_000_000,
|
||||
};
|
||||
let _pred = adapter.predict(&fv).expect("predict failed");
|
||||
}
|
||||
|
||||
// Now make two predictions with the same input (buffer is full and stable)
|
||||
let fv = FeatureVector {
|
||||
values: vec![0.1; 51],
|
||||
timestamp: 1700000000_000_000,
|
||||
};
|
||||
let pred1 = adapter.predict(&fv).expect("predict1 failed");
|
||||
let pred2 = adapter.predict(&fv).expect("predict2 failed");
|
||||
|
||||
assert!(
|
||||
(pred1.direction - pred2.direction).abs() < 1e-6,
|
||||
"Deterministic predictions should have same direction: {} vs {}",
|
||||
pred1.direction,
|
||||
pred2.direction
|
||||
);
|
||||
assert!(
|
||||
(pred1.confidence - pred2.confidence).abs() < 1e-6,
|
||||
"Deterministic predictions should have same confidence: {} vs {}",
|
||||
pred1.confidence,
|
||||
pred2.confidence
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,11 @@
|
||||
//! Per-model inference adapters for ensemble prediction
|
||||
|
||||
pub mod dqn;
|
||||
pub mod mamba2;
|
||||
pub mod ppo;
|
||||
pub mod tft;
|
||||
|
||||
pub use dqn::DqnInferenceAdapter;
|
||||
pub use mamba2::Mamba2InferenceAdapter;
|
||||
pub use ppo::PpoInferenceAdapter;
|
||||
pub use tft::TftInferenceAdapter;
|
||||
|
||||
483
ml/src/ensemble/adapters/tft.rs
Normal file
483
ml/src/ensemble/adapters/tft.rs
Normal file
@@ -0,0 +1,483 @@
|
||||
//! TFT inference adapter for ensemble prediction
|
||||
//!
|
||||
//! Wraps a Temporal Fusion Transformer model and normalizes its quantile
|
||||
//! output into a directional signal + confidence for ensemble aggregation.
|
||||
//!
|
||||
//! Unlike single-step models (DQN, PPO), TFT requires a sequence of
|
||||
//! feature vectors before it can produce a prediction. This adapter
|
||||
//! maintains an internal ring buffer that collects `sequence_length`
|
||||
//! observations before running the first forward pass.
|
||||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use candle_core::{Device, Tensor};
|
||||
|
||||
use crate::ensemble::inference_adapter::{
|
||||
EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
|
||||
};
|
||||
use crate::tft::{TFTConfig, TemporalFusionTransformer};
|
||||
use crate::{MLError, MLResult};
|
||||
|
||||
/// Inference adapter that wraps a TFT model for ensemble prediction.
|
||||
///
|
||||
/// Buffers incoming feature vectors until `sequence_length` observations
|
||||
/// are available, then runs the full TFT forward pass. The median
|
||||
/// quantile is mapped to a directional signal and the IQR provides a
|
||||
/// confidence estimate.
|
||||
#[allow(missing_debug_implementations)]
|
||||
pub struct TftInferenceAdapter {
|
||||
model: Mutex<TemporalFusionTransformer>,
|
||||
buffer: Mutex<VecDeque<FeatureVector>>,
|
||||
sequence_length: usize,
|
||||
num_static: usize,
|
||||
num_known: usize,
|
||||
num_unknown: usize,
|
||||
prediction_horizon: usize,
|
||||
num_quantiles: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
// SAFETY: TemporalFusionTransformer internally uses candle tensors which
|
||||
// are Send+Sync. The Mutex provides exclusive access for inference calls.
|
||||
unsafe impl Send for TftInferenceAdapter {}
|
||||
unsafe impl Sync for TftInferenceAdapter {}
|
||||
|
||||
impl TftInferenceAdapter {
|
||||
/// Create a new TFT inference adapter from configuration.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `config` - TFT model configuration. `input_dim` **must** equal
|
||||
/// `num_static_features + num_known_features + num_unknown_features`.
|
||||
/// * `sequence_length` - Number of feature vectors to buffer before
|
||||
/// the first forward pass. Overrides `config.sequence_length` for
|
||||
/// the buffer size (the config value is still used internally by
|
||||
/// the model).
|
||||
pub fn new(config: TFTConfig, sequence_length: usize) -> MLResult<Self> {
|
||||
let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
|
||||
let num_static = config.num_static_features;
|
||||
let num_known = config.num_known_features;
|
||||
let num_unknown = config.num_unknown_features;
|
||||
let prediction_horizon = config.prediction_horizon;
|
||||
let num_quantiles = config.num_quantiles;
|
||||
|
||||
let model = TemporalFusionTransformer::new_with_device(config, device.clone())?;
|
||||
|
||||
Ok(Self {
|
||||
model: Mutex::new(model),
|
||||
buffer: Mutex::new(VecDeque::with_capacity(sequence_length + 1)),
|
||||
sequence_length,
|
||||
num_static,
|
||||
num_known,
|
||||
num_unknown,
|
||||
prediction_horizon,
|
||||
num_quantiles,
|
||||
device,
|
||||
})
|
||||
}
|
||||
|
||||
/// Extract a fixed-length f32 slice from a feature vector starting at
|
||||
/// `offset`, zero-padding if the source is too short.
|
||||
fn extract_features(values: &[f64], offset: usize, count: usize) -> Vec<f32> {
|
||||
let mut out = vec![0.0f32; count];
|
||||
for i in 0..count {
|
||||
if let Some(&v) = values.get(offset.wrapping_add(i)) {
|
||||
if let Some(slot) = out.get_mut(i) {
|
||||
*slot = v as f32;
|
||||
}
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Build calendar-based future features for `prediction_horizon` steps.
|
||||
///
|
||||
/// Layout per step (zero-padded to `num_known`):
|
||||
/// 0: hour_sin, 1: hour_cos
|
||||
/// 2: day_sin, 3: day_cos
|
||||
/// 4: minute_sin, 5: minute_cos
|
||||
fn build_future_features(&self, base_timestamp: i64) -> Vec<f32> {
|
||||
let mut out = Vec::with_capacity(self.prediction_horizon * self.num_known);
|
||||
for step in 0..self.prediction_horizon {
|
||||
// Advance timestamp by step (assume 1-second ticks)
|
||||
let ts = base_timestamp.wrapping_add(step as i64);
|
||||
|
||||
// Derive simple cyclical features from timestamp
|
||||
let seconds_in_day = 86400.0_f64;
|
||||
let seconds_in_hour = 3600.0_f64;
|
||||
let seconds_in_minute = 60.0_f64;
|
||||
|
||||
let day_frac = (ts as f64 % seconds_in_day) / seconds_in_day;
|
||||
let hour_frac = (ts as f64 % seconds_in_hour) / seconds_in_hour;
|
||||
let minute_frac = (ts as f64 % seconds_in_minute) / seconds_in_minute;
|
||||
|
||||
let two_pi = std::f64::consts::TAU;
|
||||
|
||||
let mut step_feats = vec![0.0f32; self.num_known];
|
||||
// hour sin/cos
|
||||
if let Some(slot) = step_feats.get_mut(0) {
|
||||
*slot = (hour_frac * two_pi).sin() as f32;
|
||||
}
|
||||
if let Some(slot) = step_feats.get_mut(1) {
|
||||
*slot = (hour_frac * two_pi).cos() as f32;
|
||||
}
|
||||
// day sin/cos
|
||||
if let Some(slot) = step_feats.get_mut(2) {
|
||||
*slot = (day_frac * two_pi).sin() as f32;
|
||||
}
|
||||
if let Some(slot) = step_feats.get_mut(3) {
|
||||
*slot = (day_frac * two_pi).cos() as f32;
|
||||
}
|
||||
// minute sin/cos
|
||||
if let Some(slot) = step_feats.get_mut(4) {
|
||||
*slot = (minute_frac * two_pi).sin() as f32;
|
||||
}
|
||||
if let Some(slot) = step_feats.get_mut(5) {
|
||||
*slot = (minute_frac * two_pi).cos() as f32;
|
||||
}
|
||||
|
||||
out.extend_from_slice(&step_feats);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Map raw median quantile value to a directional signal in [-1, 1]
|
||||
/// using a smooth saturating function: sign(x) * (1 - exp(-|x|)).
|
||||
fn median_to_direction(median: f64) -> f64 {
|
||||
if median == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let sign = if median > 0.0 { 1.0 } else { -1.0 };
|
||||
let magnitude = 1.0 - (-median.abs()).exp();
|
||||
(sign * magnitude).clamp(-1.0, 1.0)
|
||||
}
|
||||
|
||||
/// Map IQR relative to median magnitude to a confidence score [0, 1].
|
||||
///
|
||||
/// Tight quantile spread (low IQR) relative to a strong signal yields
|
||||
/// high confidence. When the median is near zero the model is
|
||||
/// indecisive, so confidence should be low.
|
||||
fn iqr_to_confidence(median: f64, iqr: f64) -> f64 {
|
||||
let abs_median = median.abs();
|
||||
if abs_median < 1e-10 {
|
||||
return 0.0;
|
||||
}
|
||||
// confidence = median_strength * certainty
|
||||
// certainty = 1 / (1 + iqr / |median|)
|
||||
let certainty = 1.0 / (1.0 + iqr / abs_median);
|
||||
let strength = 1.0 - (-abs_median).exp();
|
||||
(strength * certainty).clamp(0.0, 1.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl ModelInferenceAdapter for TftInferenceAdapter {
|
||||
fn model_name(&self) -> &str {
|
||||
"TFT"
|
||||
}
|
||||
|
||||
fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// --- buffer phase (short-lived lock) ---
|
||||
let buffered_data: Option<(Vec<FeatureVector>, i64)> = {
|
||||
let mut buf = self
|
||||
.buffer
|
||||
.lock()
|
||||
.map_err(|e| MLError::LockError(format!("TFT buffer lock poisoned: {e}")))?;
|
||||
|
||||
buf.push_back(features.clone());
|
||||
|
||||
// Trim to sequence_length (keep most recent)
|
||||
while buf.len() > self.sequence_length {
|
||||
buf.pop_front();
|
||||
}
|
||||
|
||||
if buf.len() < self.sequence_length {
|
||||
None
|
||||
} else {
|
||||
// Snapshot the buffer contents so we can drop the lock
|
||||
let snapshot: Vec<FeatureVector> = buf.iter().cloned().collect();
|
||||
let last_ts = snapshot
|
||||
.last()
|
||||
.map(|fv| fv.timestamp)
|
||||
.unwrap_or(features.timestamp);
|
||||
Some((snapshot, last_ts))
|
||||
}
|
||||
}; // buffer lock dropped here
|
||||
|
||||
// If we don't have enough data yet, return a neutral prediction
|
||||
let (snapshot, last_ts) = match buffered_data {
|
||||
Some(data) => data,
|
||||
None => {
|
||||
let latency_us = start.elapsed().as_micros() as u64;
|
||||
return Ok(EnsemblePrediction {
|
||||
model_name: "TFT".to_string(),
|
||||
direction: 0.0,
|
||||
confidence: 0.0,
|
||||
metadata: PredictionMeta {
|
||||
latency_us,
|
||||
quantiles: None,
|
||||
attention_weights: None,
|
||||
q_values: None,
|
||||
},
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
// --- tensor construction ---
|
||||
|
||||
// Static features: extract from the first FV at offset 40
|
||||
let static_offset = 40;
|
||||
let first_fv_values = snapshot
|
||||
.first()
|
||||
.map(|fv| fv.values.as_slice())
|
||||
.ok_or_else(|| MLError::InferenceError("Empty snapshot".to_string()))?;
|
||||
let static_f32 = Self::extract_features(first_fv_values, static_offset, self.num_static);
|
||||
|
||||
// Historical (unknown) features: from each FV at offset 0
|
||||
let mut hist_f32 = Vec::with_capacity(self.sequence_length * self.num_unknown);
|
||||
for fv in &snapshot {
|
||||
let row = Self::extract_features(&fv.values, 0, self.num_unknown);
|
||||
hist_f32.extend_from_slice(&row);
|
||||
}
|
||||
|
||||
// Future (known) features: calendar-derived
|
||||
let future_f32 = self.build_future_features(last_ts);
|
||||
|
||||
// Build tensors
|
||||
let static_tensor =
|
||||
Tensor::from_vec(static_f32, (1, self.num_static), &self.device).map_err(|e| {
|
||||
MLError::ModelError(format!("Failed to create static tensor: {e}"))
|
||||
})?;
|
||||
|
||||
let hist_tensor = Tensor::from_vec(
|
||||
hist_f32,
|
||||
(1, self.sequence_length, self.num_unknown),
|
||||
&self.device,
|
||||
)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to create historical tensor: {e}")))?;
|
||||
|
||||
let future_tensor = Tensor::from_vec(
|
||||
future_f32,
|
||||
(1, self.prediction_horizon, self.num_known),
|
||||
&self.device,
|
||||
)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to create future tensor: {e}")))?;
|
||||
|
||||
// --- forward pass (model lock) ---
|
||||
let quantile_preds = {
|
||||
let mut model = self
|
||||
.model
|
||||
.lock()
|
||||
.map_err(|e| MLError::LockError(format!("TFT model lock poisoned: {e}")))?;
|
||||
model.forward(&static_tensor, &hist_tensor, &future_tensor)?
|
||||
};
|
||||
|
||||
// --- output processing ---
|
||||
// quantile_preds shape: [1, prediction_horizon, num_quantiles]
|
||||
let squeezed = quantile_preds
|
||||
.squeeze(0)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to squeeze TFT output: {e}")))?;
|
||||
let pred_data: Vec<Vec<f32>> = squeezed
|
||||
.to_vec2()
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to extract TFT output: {e}")))?;
|
||||
|
||||
// Use first horizon step for directional signal
|
||||
let first_horizon = pred_data
|
||||
.first()
|
||||
.ok_or_else(|| MLError::InferenceError("TFT produced empty output".to_string()))?;
|
||||
|
||||
let median_idx = self.num_quantiles / 2;
|
||||
let q25_idx = self.num_quantiles / 4;
|
||||
let q75_idx = (self.num_quantiles * 3) / 4;
|
||||
|
||||
let median = first_horizon
|
||||
.get(median_idx)
|
||||
.copied()
|
||||
.unwrap_or(0.0) as f64;
|
||||
let q25 = first_horizon
|
||||
.get(q25_idx)
|
||||
.copied()
|
||||
.unwrap_or(0.0) as f64;
|
||||
let q75 = first_horizon
|
||||
.get(q75_idx)
|
||||
.copied()
|
||||
.unwrap_or(0.0) as f64;
|
||||
let iqr = (q75 - q25).abs();
|
||||
|
||||
let direction = Self::median_to_direction(median);
|
||||
let confidence = Self::iqr_to_confidence(median, iqr);
|
||||
|
||||
// Collect all quantiles for metadata (first horizon)
|
||||
let quantile_values: Vec<f64> = first_horizon.iter().map(|&v| v as f64).collect();
|
||||
|
||||
let latency_us = start.elapsed().as_micros() as u64;
|
||||
|
||||
Ok(EnsemblePrediction {
|
||||
model_name: "TFT".to_string(),
|
||||
direction,
|
||||
confidence,
|
||||
metadata: PredictionMeta {
|
||||
latency_us,
|
||||
quantiles: Some(quantile_values),
|
||||
attention_weights: None,
|
||||
q_values: None,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn is_ready(&self) -> bool {
|
||||
let buf_ok = self
|
||||
.buffer
|
||||
.lock()
|
||||
.map(|buf| buf.len() >= self.sequence_length)
|
||||
.unwrap_or(false);
|
||||
let model_ok = self.model.lock().is_ok();
|
||||
buf_ok && model_ok
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Small config for fast tests.
|
||||
/// input_dim = num_static + num_known + num_unknown = 6 + 6 + 8 = 20
|
||||
fn test_config() -> TFTConfig {
|
||||
TFTConfig {
|
||||
input_dim: 20,
|
||||
hidden_dim: 32,
|
||||
num_heads: 2,
|
||||
num_layers: 1,
|
||||
prediction_horizon: 5,
|
||||
sequence_length: 4,
|
||||
num_quantiles: 9,
|
||||
num_static_features: 6,
|
||||
num_known_features: 6,
|
||||
num_unknown_features: 8,
|
||||
dropout_rate: 0.0,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_fv(ts: i64) -> FeatureVector {
|
||||
// 51 values — more than our small config needs, adapter will
|
||||
// extract/zero-pad as required.
|
||||
FeatureVector {
|
||||
values: vec![0.1; 51],
|
||||
timestamp: ts,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tft_adapter_creation() {
|
||||
let adapter = TftInferenceAdapter::new(test_config(), 4)
|
||||
.expect("TftInferenceAdapter::new should succeed");
|
||||
assert_eq!(adapter.model_name(), "TFT");
|
||||
// Buffer is empty, so the adapter is NOT ready yet
|
||||
assert!(
|
||||
!adapter.is_ready(),
|
||||
"Adapter should not be ready with an empty buffer"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tft_adapter_buffers_and_predicts() {
|
||||
let seq_len = 4;
|
||||
let adapter = TftInferenceAdapter::new(test_config(), seq_len)
|
||||
.expect("TftInferenceAdapter::new should succeed");
|
||||
|
||||
// Feed seq_len - 1 feature vectors; each should return neutral
|
||||
for i in 0..(seq_len - 1) {
|
||||
let fv = make_fv(1_700_000_000 + i as i64);
|
||||
let pred = adapter.predict(&fv).expect("predict should not error");
|
||||
assert_eq!(
|
||||
pred.direction, 0.0,
|
||||
"Should return neutral direction while buffering"
|
||||
);
|
||||
assert_eq!(
|
||||
pred.confidence, 0.0,
|
||||
"Should return zero confidence while buffering"
|
||||
);
|
||||
}
|
||||
|
||||
assert!(
|
||||
!adapter.is_ready(),
|
||||
"Adapter should still not be ready with {} / {} entries",
|
||||
seq_len - 1,
|
||||
seq_len
|
||||
);
|
||||
|
||||
// Feed the final FV to fill the buffer
|
||||
let fv_final = make_fv(1_700_000_000 + seq_len as i64);
|
||||
let pred = adapter
|
||||
.predict(&fv_final)
|
||||
.expect("predict should succeed once buffer is full");
|
||||
|
||||
// Now the adapter should be ready
|
||||
assert!(
|
||||
adapter.is_ready(),
|
||||
"Adapter should be ready after receiving seq_len entries"
|
||||
);
|
||||
|
||||
// Direction and confidence should be bounded
|
||||
assert!(
|
||||
pred.direction >= -1.0 && pred.direction <= 1.0,
|
||||
"direction {} out of [-1,1]",
|
||||
pred.direction
|
||||
);
|
||||
assert!(
|
||||
pred.confidence >= 0.0 && pred.confidence <= 1.0,
|
||||
"confidence {} out of [0,1]",
|
||||
pred.confidence
|
||||
);
|
||||
|
||||
// Quantile metadata should be present
|
||||
assert!(
|
||||
pred.metadata.quantiles.is_some(),
|
||||
"TFT predictions should include quantile metadata"
|
||||
);
|
||||
|
||||
// One more prediction to ensure repeated calls work
|
||||
let fv_extra = make_fv(1_700_000_000 + (seq_len + 1) as i64);
|
||||
let pred2 = adapter
|
||||
.predict(&fv_extra)
|
||||
.expect("subsequent predict should succeed");
|
||||
assert!(
|
||||
pred2.direction >= -1.0 && pred2.direction <= 1.0,
|
||||
"second direction {} out of [-1,1]",
|
||||
pred2.direction
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tft_adapter_deterministic() {
|
||||
let seq_len = 4;
|
||||
// Create two adapters with identical configs
|
||||
let adapter1 = TftInferenceAdapter::new(test_config(), seq_len)
|
||||
.expect("adapter1 creation");
|
||||
let adapter2 = TftInferenceAdapter::new(test_config(), seq_len)
|
||||
.expect("adapter2 creation");
|
||||
|
||||
// Feed identical sequences
|
||||
for i in 0..seq_len {
|
||||
let fv = make_fv(1_700_000_000 + i as i64);
|
||||
let _ = adapter1.predict(&fv);
|
||||
let _ = adapter2.predict(&fv);
|
||||
}
|
||||
|
||||
// The final prediction after filling the buffer should be identical
|
||||
let fv = make_fv(1_700_000_000 + seq_len as i64);
|
||||
let pred1 = adapter1.predict(&fv).expect("adapter1 predict");
|
||||
let pred2 = adapter2.predict(&fv).expect("adapter2 predict");
|
||||
|
||||
assert_eq!(
|
||||
pred1.direction, pred2.direction,
|
||||
"Deterministic predictions should have same direction"
|
||||
);
|
||||
assert_eq!(
|
||||
pred1.confidence, pred2.confidence,
|
||||
"Deterministic predictions should have same confidence"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user