From c402da4c8b5a4bfef100c950033622190e4c13fe Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 21 Feb 2026 13:27:11 +0100 Subject: [PATCH] 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 --- ml/src/ensemble/adapters/mamba2.rs | 319 +++++++++++++++++++ ml/src/ensemble/adapters/mod.rs | 4 + ml/src/ensemble/adapters/tft.rs | 483 +++++++++++++++++++++++++++++ 3 files changed, 806 insertions(+) create mode 100644 ml/src/ensemble/adapters/mamba2.rs create mode 100644 ml/src/ensemble/adapters/tft.rs diff --git a/ml/src/ensemble/adapters/mamba2.rs b/ml/src/ensemble/adapters/mamba2.rs new file mode 100644 index 000000000..ba7deff71 --- /dev/null +++ b/ml/src/ensemble/adapters/mamba2.rs @@ -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, + buffer: Mutex>>, + 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 { + 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 { + 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 { + 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 = 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 + ); + } +} diff --git a/ml/src/ensemble/adapters/mod.rs b/ml/src/ensemble/adapters/mod.rs index e653181ab..4e1627084 100644 --- a/ml/src/ensemble/adapters/mod.rs +++ b/ml/src/ensemble/adapters/mod.rs @@ -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; diff --git a/ml/src/ensemble/adapters/tft.rs b/ml/src/ensemble/adapters/tft.rs new file mode 100644 index 000000000..5e6a51325 --- /dev/null +++ b/ml/src/ensemble/adapters/tft.rs @@ -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, + buffer: Mutex>, + 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 { + 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 { + 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 { + 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 { + let start = std::time::Instant::now(); + + // --- buffer phase (short-lived lock) --- + let buffered_data: Option<(Vec, 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 = 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> = 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 = 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" + ); + } +}