foxhunt/crates/ml-core/src/training.rs

//! Unified Training Trait for All ML Models
//!
//! Provides a common interface for training all models (MAMBA-2, DQN, PPO, TFT, etc.).
//! Standardizes batch processing, gradient computation, optimizer steps, checkpointing,
//! and metrics collection across all model types.
//!
//! Uses `GpuTensor` and `MlDevice` instead of candle types.

use crate::MLError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Training metrics collected during training
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingMetrics {
    /// Training loss
    pub loss: f64,
    /// Validation loss
    pub val_loss: Option<f64>,
    /// Accuracy metric (model-specific)
    pub accuracy: Option<f64>,
    /// Learning rate used in this step
    pub learning_rate: f64,
    /// Gradient norm (for monitoring gradient explosion)
    pub grad_norm: Option<f64>,
    /// Custom model-specific metrics
    pub custom_metrics: HashMap<String, f64>,
}

impl Default for TrainingMetrics {
    fn default() -> Self {
        Self {
            loss: 0.0,
            val_loss: None,
            accuracy: None,
            learning_rate: 1e-4,
            grad_norm: None,
            custom_metrics: HashMap::new(),
        }
    }
}

/// Checkpoint metadata for standardized format
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CheckpointMetadata {
    /// Model type (MAMBA-2, DQN, PPO, TFT)
    pub model_type: String,
    /// Model version
    pub version: String,
    /// Epoch number
    pub epoch: usize,
    /// Training step
    pub step: usize,
    /// Timestamp when checkpoint was created
    pub timestamp: std::time::SystemTime,
    /// Training configuration (serialized as JSON)
    pub config: serde_json::Value,
    /// Performance metrics at checkpoint time
    pub metrics: TrainingMetrics,
}

/// Unified training interface for all ML models
///
/// Implementations must provide:
/// - Forward pass (inference)
/// - Backward pass (gradient computation)
/// - Optimizer step (parameter updates)
/// - Checkpoint save/load (safetensors + JSON metadata)
/// - Metrics collection (loss, accuracy, custom metrics)
///
/// This trait uses opaque types (f64 for loss, &str for device info) so that
/// it does not depend on any specific tensor library.
pub trait UnifiedTrainable {
    /// Get model type identifier (MAMBA-2, DQN, PPO, TFT)
    fn model_type(&self) -> &str;

    /// Get device description (e.g., "cpu", "cuda:0")
    fn device_name(&self) -> String;

    /// Forward pass through model (returns loss as f64)
    fn forward_loss(&mut self, input: &[f32], target: &[f32]) -> Result<f64, MLError>;

    /// Backward pass to compute gradients
    fn backward(&mut self, loss_value: f64) -> Result<f64, MLError>;

    /// Update model parameters using optimizer
    fn optimizer_step(&mut self) -> Result<(), MLError>;

    /// Zero gradients before next backward pass
    fn zero_grad(&mut self) -> Result<(), MLError>;

    /// Get current learning rate
    fn get_learning_rate(&self) -> f64;

    /// Set learning rate (for scheduling)
    fn set_learning_rate(&mut self, lr: f64) -> Result<(), MLError>;

    /// Get current training step count
    fn get_step(&self) -> usize;

    /// Collect current training metrics
    fn collect_metrics(&self) -> TrainingMetrics;

    /// Save model checkpoint in standardized format
    ///
    /// Format: safetensors for weights, JSON for metadata
    fn save_checkpoint(&self, checkpoint_path: &str) -> Result<String, MLError>;

    /// Load model checkpoint from standardized format
    fn load_checkpoint(&mut self, checkpoint_path: &str) -> Result<CheckpointMetadata, MLError>;
}

/// Helper functions for standardized checkpoint management
pub mod checkpoint {
    use super::*;
    use std::path::Path;

    /// Generate checkpoint filename with epoch and step
    pub fn checkpoint_filename(model_type: &str, epoch: usize, step: usize) -> String {
        format!("{}_epoch{}_step{}", model_type, epoch, step)
    }

    /// Save checkpoint metadata to JSON
    pub fn save_metadata(
        metadata: &CheckpointMetadata,
        checkpoint_path: &str,
    ) -> Result<(), MLError> {
        let metadata_path = format!("{}.json", checkpoint_path);
        let json = serde_json::to_string_pretty(metadata).map_err(|e| {
            MLError::ModelError(format!("Failed to serialize checkpoint metadata: {}", e))
        })?;

        std::fs::write(&metadata_path, json).map_err(|e| {
            MLError::ModelError(format!("Failed to write checkpoint metadata: {}", e))
        })?;

        Ok(())
    }

    /// Load checkpoint metadata from JSON
    pub fn load_metadata(checkpoint_path: &str) -> Result<CheckpointMetadata, MLError> {
        let metadata_path = format!("{}.json", checkpoint_path);

        if !Path::new(&metadata_path).exists() {
            return Err(MLError::ModelError(format!(
                "Checkpoint metadata not found: {}",
                metadata_path
            )));
        }

        let json = std::fs::read_to_string(&metadata_path).map_err(|e| {
            MLError::ModelError(format!("Failed to read checkpoint metadata: {}", e))
        })?;

        let metadata: CheckpointMetadata = serde_json::from_str(&json).map_err(|e| {
            MLError::ModelError(format!("Failed to deserialize checkpoint metadata: {}", e))
        })?;

        Ok(metadata)
    }

    /// Check if checkpoint exists
    pub fn checkpoint_exists(checkpoint_path: &str) -> bool {
        let safetensors_path = format!("{}.safetensors", checkpoint_path);
        let metadata_path = format!("{}.json", checkpoint_path);

        Path::new(&safetensors_path).exists() && Path::new(&metadata_path).exists()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_training_metrics_default() {
        let metrics = TrainingMetrics::default();
        assert_eq!(metrics.loss, 0.0);
        assert_eq!(metrics.val_loss, None);
        assert_eq!(metrics.learning_rate, 1e-4);
    }

    #[test]
    fn test_checkpoint_filename() {
        let filename = checkpoint::checkpoint_filename("MAMBA-2", 10, 1000);
        assert_eq!(filename, "MAMBA-2_epoch10_step1000");
    }

    #[test]
    fn test_checkpoint_metadata_serialization() -> anyhow::Result<()> {
        let metadata = CheckpointMetadata {
            model_type: "MAMBA-2".to_owned(),
            version: "2.0.0".to_owned(),
            epoch: 10,
            step: 1000,
            timestamp: std::time::SystemTime::now(),
            config: serde_json::json!({"d_model": 256}),
            metrics: TrainingMetrics::default(),
        };

        let json = serde_json::to_string_pretty(&metadata)?;
        assert!(json.contains("MAMBA-2"));

        let deserialized: CheckpointMetadata = serde_json::from_str(&json)?;
        assert_eq!(deserialized.model_type, "MAMBA-2");
        assert_eq!(deserialized.epoch, 10);

        Ok(())
    }
}