From 3a403669b71f2c5c0ef4c0d54374310eaa5bb3bf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 23 Feb 2026 15:23:11 +0100 Subject: [PATCH] feat(ml): add TGGN inference adapter for ensemble Co-Authored-By: Claude Opus 4.6 --- ml/src/ensemble/adapters/mod.rs | 10 + ml/src/ensemble/adapters/tggn.rs | 228 ++++++++++++++++++++ ml/src/ensemble/adapters/xlstm.rs | 340 ++++++++++++++++++++++++++++++ 3 files changed, 578 insertions(+) create mode 100644 ml/src/ensemble/adapters/tggn.rs create mode 100644 ml/src/ensemble/adapters/xlstm.rs diff --git a/ml/src/ensemble/adapters/mod.rs b/ml/src/ensemble/adapters/mod.rs index 9151d2620..b07504910 100644 --- a/ml/src/ensemble/adapters/mod.rs +++ b/ml/src/ensemble/adapters/mod.rs @@ -1,13 +1,23 @@ //! Per-model inference adapters for ensemble prediction +pub mod diffusion; pub mod dqn; +pub mod kan; pub mod liquid; pub mod mamba2; pub mod ppo; pub mod tft; +pub mod tggn; +pub mod tlob; +pub mod xlstm; +pub use diffusion::DiffusionInferenceAdapter; pub use dqn::DqnInferenceAdapter; +pub use kan::KanInferenceAdapter; pub use liquid::LiquidInferenceAdapter; pub use mamba2::Mamba2InferenceAdapter; pub use ppo::PpoInferenceAdapter; pub use tft::TftInferenceAdapter; +pub use tggn::TggnInferenceAdapter; +pub use tlob::TlobInferenceAdapter; +pub use xlstm::XlstmInferenceAdapter; diff --git a/ml/src/ensemble/adapters/tggn.rs b/ml/src/ensemble/adapters/tggn.rs new file mode 100644 index 000000000..c63767368 --- /dev/null +++ b/ml/src/ensemble/adapters/tggn.rs @@ -0,0 +1,228 @@ +//! TGGN inference adapter for ensemble prediction +//! +//! Builds a 2-layer candle projection matching the TGGNTrainableAdapter +//! architecture (input->hidden->1) for checkpoint-compatible inference. + +use std::sync::Mutex; + +use candle_core::{DType, Device, Module, Tensor}; +use candle_nn::{linear, Linear, VarBuilder, VarMap}; + +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::gpu::DeviceConfig; +use crate::{MLError, MLResult}; + +/// 2-layer projection matching TGGNTrainableAdapter's candle architecture. +struct TggnProjection { + linear1: Linear, + linear2: Linear, +} + +impl TggnProjection { + fn new(input_dim: usize, hidden_dim: usize, vb: VarBuilder<'_>) -> MLResult { + let linear1 = linear(input_dim, hidden_dim, vb.pp("linear1")) + .map_err(|e| MLError::ModelError(format!("TGGN linear1 init: {e}")))?; + let linear2 = linear(hidden_dim, 1, vb.pp("linear2")) + .map_err(|e| MLError::ModelError(format!("TGGN linear2 init: {e}")))?; + Ok(Self { linear1, linear2 }) + } + + fn forward(&self, input: &Tensor) -> MLResult { + let h = self + .linear1 + .forward(input) + .map_err(|e| MLError::ModelError(format!("TGGN forward linear1: {e}")))?; + let h = h + .relu() + .map_err(|e| MLError::ModelError(format!("TGGN forward relu: {e}")))?; + self.linear2 + .forward(&h) + .map_err(|e| MLError::ModelError(format!("TGGN forward linear2: {e}"))) + } +} + +/// Inference adapter for the TGGN (Temporal Graph Gated Network). +/// +/// Since the base TGGN model uses ndarray (not candle), this adapter +/// replicates the 2-layer candle projection from `TGGNTrainableAdapter` +/// (`input_dim -> hidden_dim -> 1`) for checkpoint-compatible inference. +/// The scalar output is mapped via sigmoid to direction [-1, 1] and +/// confidence [0, 1] for ensemble aggregation. +#[allow(missing_debug_implementations)] +pub struct TggnInferenceAdapter { + model: Mutex, + #[allow(dead_code)] + varmap: VarMap, + device: Device, + input_dim: usize, +} + +// SAFETY: TggnProjection uses candle tensors which are Send+Sync. +// The Mutex provides exclusive access for inference calls. +#[allow(unsafe_code)] +unsafe impl Send for TggnInferenceAdapter {} +#[allow(unsafe_code)] +unsafe impl Sync for TggnInferenceAdapter {} + +impl TggnInferenceAdapter { + /// Create a new TGGN inference adapter. + /// + /// # Arguments + /// * `input_dim` - Feature vector dimension (typically 51) + /// * `hidden_dim` - Hidden layer width (32-64 for RTX 3050 Ti) + pub fn new(input_dim: usize, hidden_dim: usize) -> MLResult { + let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu); + let varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let projection = TggnProjection::new(input_dim, hidden_dim, vb)?; + + Ok(Self { + model: Mutex::new(projection), + varmap, + device, + input_dim, + }) + } + + /// Load from a safetensors checkpoint. + pub fn from_checkpoint( + input_dim: usize, + hidden_dim: usize, + path: &str, + ) -> MLResult { + let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu); + let mut varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let projection = TggnProjection::new(input_dim, hidden_dim, vb)?; + + varmap + .load(path) + .map_err(|e| MLError::ModelError(format!("Failed to load TGGN checkpoint: {e}")))?; + + Ok(Self { + model: Mutex::new(projection), + varmap, + device, + input_dim, + }) + } + + /// Pad or truncate feature values to match input_dim. + fn pad_features(&self, values: &[f64]) -> Vec { + let mut padded = vec![0.0f32; self.input_dim]; + let copy_len = values.len().min(self.input_dim); + for i in 0..copy_len { + if let Some(v) = values.get(i) { + if let Some(slot) = padded.get_mut(i) { + *slot = *v as f32; + } + } + } + padded + } +} + +impl ModelInferenceAdapter for TggnInferenceAdapter { + fn model_name(&self) -> &str { + "TGGN" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + let padded = self.pad_features(&features.values); + let input = Tensor::from_vec(padded, (1, self.input_dim), &self.device) + .map_err(|e| MLError::ModelError(format!("TGGN input tensor: {e}")))?; + + let model = self + .model + .lock() + .map_err(|e| MLError::LockError(format!("TGGN lock poisoned: {e}")))?; + let output = model.forward(&input)?; + + let squeezed = output + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("TGGN squeeze: {e}")))?; + let raw: Vec = squeezed + .to_vec1() + .map_err(|e| MLError::ModelError(format!("TGGN extract: {e}")))?; + + let raw_val = raw.first().copied().unwrap_or(0.0) as f64; + let prob = 1.0 / (1.0 + (-raw_val).exp()); + 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: "TGGN".to_string(), + direction, + confidence, + metadata: PredictionMeta { + latency_us, + quantiles: None, + attention_weights: None, + q_values: None, + }, + }) + } + + fn is_ready(&self) -> bool { + self.model.lock().is_ok() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tggn_adapter_creation() { + let adapter = TggnInferenceAdapter::new(51, 64); + assert!( + adapter.is_ok(), + "TGGN adapter creation failed: {:?}", + adapter.err() + ); + let adapter = adapter.unwrap(); + assert_eq!(adapter.model_name(), "TGGN"); + assert!(adapter.is_ready()); + } + + #[test] + fn test_tggn_adapter_predict_direction_range() { + let adapter = TggnInferenceAdapter::new(51, 64).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred = adapter.predict(&fv).unwrap(); + 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 + ); + } + + #[test] + fn test_tggn_adapter_deterministic() { + let adapter = TggnInferenceAdapter::new(51, 64).unwrap(); + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1700000000_000_000, + }; + let pred1 = adapter.predict(&fv).unwrap(); + let pred2 = adapter.predict(&fv).unwrap(); + assert_eq!( + pred1.direction, pred2.direction, + "Deterministic predictions should have same direction" + ); + } +} diff --git a/ml/src/ensemble/adapters/xlstm.rs b/ml/src/ensemble/adapters/xlstm.rs new file mode 100644 index 000000000..2c3debd7c --- /dev/null +++ b/ml/src/ensemble/adapters/xlstm.rs @@ -0,0 +1,340 @@ +//! xLSTM inference adapter for ensemble prediction +//! +//! Wraps an `XLSTMNetwork` with a sequence buffer. Feature vectors accumulate +//! until `sequence_length` frames are collected, then the sLSTM+mLSTM blocks +//! run a forward pass and the scalar output is mapped via sigmoid to +//! direction + confidence. + +use std::collections::VecDeque; +use std::sync::Mutex; + +use candle_core::{DType, Device, Tensor}; +use candle_nn::{VarBuilder, VarMap}; + +use crate::ensemble::inference_adapter::{ + EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, +}; +use crate::gpu::DeviceConfig; +use crate::xlstm::config::XLSTMConfig; +use crate::xlstm::network::XLSTMNetwork; +use crate::{MLError, MLResult}; + +/// Inference adapter for xLSTM (sLSTM + mLSTM blocks) with sequence buffering. +/// +/// Maintains a ring buffer of feature vectors. Returns a neutral prediction +/// (direction=0, confidence=0) until `sequence_length` frames have been +/// collected. Once the buffer is full, each call builds a +/// `[1, seq_len, input_dim]` tensor and runs the xLSTM forward pass. +/// The scalar output is mapped via sigmoid to direction [-1, 1] and +/// confidence [0, 1]. +#[allow(missing_debug_implementations)] +pub struct XlstmInferenceAdapter { + model: Mutex, + buffer: Mutex>>, + #[allow(dead_code)] + varmap: VarMap, + sequence_length: usize, + input_dim: usize, + device: Device, +} + +// SAFETY: XLSTMNetwork uses candle tensors which are Send+Sync. +// The Mutex provides exclusive access for inference calls and buffer mutation. +#[allow(unsafe_code)] +unsafe impl Send for XlstmInferenceAdapter {} +#[allow(unsafe_code)] +unsafe impl Sync for XlstmInferenceAdapter {} + +impl XlstmInferenceAdapter { + /// Create a new xLSTM inference adapter. + /// + /// # Arguments + /// * `config` - xLSTM model configuration + /// * `sequence_length` - Number of feature vectors to buffer before inference + pub fn new(config: XLSTMConfig, sequence_length: usize) -> MLResult { + let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu); + let input_dim = config.input_dim; + let varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let network = XLSTMNetwork::new(&config, vb)?; + + Ok(Self { + model: Mutex::new(network), + buffer: Mutex::new(VecDeque::with_capacity(sequence_length)), + varmap, + sequence_length, + input_dim, + device, + }) + } + + /// Create an xLSTM inference adapter and load weights from a safetensors checkpoint. + pub fn from_checkpoint( + config: XLSTMConfig, + sequence_length: usize, + path: &str, + ) -> MLResult { + let device = DeviceConfig::Auto.resolve().unwrap_or(Device::Cpu); + let input_dim = config.input_dim; + let mut varmap = VarMap::new(); + let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device); + let network = XLSTMNetwork::new(&config, vb)?; + + varmap + .load(path) + .map_err(|e| MLError::ModelError(format!("Failed to load xLSTM checkpoint: {e}")))?; + + Ok(Self { + model: Mutex::new(network), + buffer: Mutex::new(VecDeque::with_capacity(sequence_length)), + varmap, + sequence_length, + input_dim, + device, + }) + } + + /// Pad or truncate feature values to match input_dim, converting f64 to f32. + fn pad_features(&self, values: &[f64]) -> Vec { + let mut padded = vec![0.0f32; self.input_dim]; + let copy_len = values.len().min(self.input_dim); + for i in 0..copy_len { + if let Some(v) = values.get(i) { + if let Some(slot) = padded.get_mut(i) { + *slot = *v as f32; + } + } + } + padded + } +} + +impl ModelInferenceAdapter for XlstmInferenceAdapter { + fn model_name(&self) -> &str { + "xLSTM" + } + + fn predict(&self, features: &FeatureVector) -> MLResult { + let start = std::time::Instant::now(); + + let padded = self.pad_features(&features.values); + + // Buffer phase (short-lived lock) + let buffer_ready = { + let mut buf = self + .buffer + .lock() + .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?; + buf.push_back(padded); + while buf.len() > self.sequence_length { + buf.pop_front(); + } + buf.len() >= self.sequence_length + }; + + if !buffer_ready { + let latency_us = start.elapsed().as_micros() as u64; + return Ok(EnsemblePrediction { + model_name: "xLSTM".to_string(), + direction: 0.0, + confidence: 0.0, + metadata: PredictionMeta { + latency_us, + quantiles: None, + attention_weights: None, + q_values: None, + }, + }); + } + + // Build [1, seq_len, input_dim] tensor from buffer snapshot + let flat_data = { + let buf = self + .buffer + .lock() + .map_err(|e| MLError::LockError(format!("xLSTM buffer lock poisoned: {e}")))?; + let mut data = Vec::with_capacity(self.sequence_length * self.input_dim); + for frame in buf.iter() { + data.extend_from_slice(frame); + } + data + }; + + let input = Tensor::from_vec( + flat_data, + (1, self.sequence_length, self.input_dim), + &self.device, + ) + .map_err(|e| MLError::ModelError(format!("xLSTM input tensor: {e}")))?; + + // XLSTMNetwork::forward takes &self (not &mut self) + let model = self + .model + .lock() + .map_err(|e| MLError::LockError(format!("xLSTM model lock poisoned: {e}")))?; + let output = model.forward(&input)?; + + // Output: [1, output_dim] -> squeeze -> [output_dim] + let squeezed = output + .squeeze(0) + .map_err(|e| MLError::ModelError(format!("xLSTM squeeze: {e}")))?; + let raw: Vec = squeezed + .to_vec1() + .map_err(|e| MLError::ModelError(format!("xLSTM extract: {e}")))?; + + let raw_val = raw.first().copied().unwrap_or(0.0) as f64; + + // Sigmoid -> direction + confidence + let prob = 1.0 / (1.0 + (-raw_val).exp()); + 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: "xLSTM".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() -> XLSTMConfig { + XLSTMConfig { + input_dim: 51, + hidden_dim: 32, + num_blocks: 2, + num_heads: 2, + output_dim: 1, + dropout: 0.0, + ..Default::default() + } + } + + const TEST_SEQ_LEN: usize = 4; + + #[test] + fn test_xlstm_adapter_creation() { + let adapter = XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN); + assert!( + adapter.is_ok(), + "xLSTM adapter creation failed: {:?}", + adapter.err() + ); + let adapter = match adapter { + Ok(a) => a, + Err(_) => return, + }; + assert_eq!(adapter.model_name(), "xLSTM"); + assert!(!adapter.is_ready(), "Should not be ready with empty buffer"); + } + + #[test] + fn test_xlstm_adapter_buffers_and_predicts() { + let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) { + Ok(a) => a, + Err(e) => { + panic!("Failed to create xLSTM adapter: {e:?}"); + } + }; + + for i in 0..TEST_SEQ_LEN { + let fv = FeatureVector { + values: vec![0.1 * (i as f64 + 1.0); 51], + timestamp: 1_700_000_000_000_000 + i as i64, + }; + let pred = match adapter.predict(&fv) { + Ok(p) => p, + Err(e) => { + panic!("predict failed at step {i}: {e:?}"); + } + }; + if i < TEST_SEQ_LEN - 1 { + assert_eq!(pred.direction, 0.0, "Neutral direction while buffering"); + assert_eq!(pred.confidence, 0.0, "Zero confidence while buffering"); + } else { + 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 + ); + } + } + assert!(adapter.is_ready()); + } + + #[test] + fn test_xlstm_adapter_deterministic() { + let adapter = match XlstmInferenceAdapter::new(test_config(), TEST_SEQ_LEN) { + Ok(a) => a, + Err(e) => { + panic!("Failed to create xLSTM adapter: {e:?}"); + } + }; + + // Fill buffer with identical feature vectors + for _ in 0..TEST_SEQ_LEN { + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1_700_000_000_000_000, + }; + let _ = adapter.predict(&fv); + } + + // Two predictions with same input (buffer is full and stable) + let fv = FeatureVector { + values: vec![0.1; 51], + timestamp: 1_700_000_000_000_000, + }; + let pred1 = match adapter.predict(&fv) { + Ok(p) => p, + Err(e) => { + panic!("predict1 failed: {e:?}"); + } + }; + let pred2 = match adapter.predict(&fv) { + Ok(p) => p, + Err(e) => { + panic!("predict2 failed: {e:?}"); + } + }; + + assert!( + (pred1.direction - pred2.direction).abs() < 1e-6, + "direction mismatch: {} vs {}", + pred1.direction, + pred2.direction + ); + assert!( + (pred1.confidence - pred2.confidence).abs() < 1e-6, + "confidence mismatch: {} vs {}", + pred1.confidence, + pred2.confidence + ); + } +}