foxhunt/ml/examples/tft_int8_calibration_simple.rs

//! Simplified TFT INT8 Calibration (No Quantized Dependencies)
//!
//! Creates calibration dataset from ES.FUT DBN data for INT8 quantization.
//! This version avoids broken quantized_tft/lstm/attention modules.

use anyhow::{Context, Result};
use candle_core::{DType, Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use tracing::{info, warn};

use ml::data_loaders::DbnSequenceLoader;
use ml::tft::{TFTConfig, TemporalFusionTransformer};

/// Per-layer quantization parameters
#[derive(Debug, Clone, Serialize, Deserialize)]
struct LayerQuantizationParams {
    scale: f32,
    zero_point: i8,
    min_val: f32,
    max_val: f32,
    num_samples: usize,
}

/// Calibration dataset
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CalibrationData {
    num_samples: usize,
    layers: HashMap<String, LayerQuantizationParams>,
    data_source: String,
    generated_at: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!("  TFT INT8 Calibration (Simplified)");
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!();

    // Load DBN data (use ES.FUT_ohlcv-1m_2024-01-02.dbn - small file)
    // Note: DBN decoder requires uncompressed .dbn files, not .dbn.zst
    let dbn_file = PathBuf::from("test_data/real/databento");
    if !dbn_file.exists() {
        return Err(anyhow::anyhow!(
            "DBN directory not found: {}",
            dbn_file.display()
        ));
    }

    // Check for ES.FUT file (small, single-day)
    let es_fut_path = dbn_file.join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
    if !es_fut_path.exists() {
        return Err(anyhow::anyhow!(
            "ES.FUT file not found: {}. Please ensure uncompressed DBN files are available.",
            es_fut_path.display()
        ));
    }

    info!("Loading ES.FUT data from: {:?}", dbn_file);
    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?;
    let (train_data, _) = loader.load_sequences(&dbn_file, 0.9).await?;
    info!("Loaded {} sequences", train_data.len());

    // Create TFT
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = TFTConfig {
        input_dim: 256,
        hidden_dim: 64,
        num_heads: 4,
        num_layers: 2,
        prediction_horizon: 10,
        sequence_length: 60,
        num_quantiles: 3,
        num_static_features: 2,
        num_known_features: 3,
        num_unknown_features: 256,
        batch_size: 1,
        ..Default::default()
    };

    let mut tft = TemporalFusionTransformer::new(config)?;
    info!("Created TFT model");

    // Run calibration
    info!("Running calibration forward passes...");
    let mut activation_stats: HashMap<String, Vec<(f32, f32)>> = HashMap::new();

    for (idx, (input, _)) in train_data.iter().take(50).enumerate() {
        if idx % 10 == 0 {
            info!("  Progress: {}/50", idx + 1);
        }

        let batch = input.dims()[0];
        let static_features = Tensor::zeros((batch, 2), DType::F32, &device)?;
        let historical_features = input.to_dtype(DType::F32)?;
        let future_features = Tensor::zeros((batch, 10, 3), DType::F32, &device)?;

        let output = tft.forward(&static_features, &historical_features, &future_features)?;

        // Collect stats
        let vec = output.flatten_all()?.to_vec1::<f32>()?;
        let min_val = vec.iter().cloned().fold(f32::INFINITY, f32::min);
        let max_val = vec.iter().cloned().fold(f32::NEG_INFINITY, f32::max);

        activation_stats
            .entry("output_layer".to_string())
            .or_default()
            .push((min_val, max_val));
    }

    // Calculate quantization parameters
    let mut layers = HashMap::new();
    for (layer_name, stats) in activation_stats {
        let global_min = stats
            .iter()
            .map(|(min, _)| *min)
            .fold(f32::INFINITY, f32::min);
        let global_max = stats
            .iter()
            .map(|(_, max)| *max)
            .fold(f32::NEG_INFINITY, f32::max);

        let abs_max = global_min.abs().max(global_max.abs());
        let scale = if abs_max > 0.0 { abs_max / 127.0 } else { 1.0 };
        let zero_point = 127i8;

        layers.insert(
            layer_name,
            LayerQuantizationParams {
                scale,
                zero_point,
                min_val: global_min,
                max_val: global_max,
                num_samples: stats.len(),
            },
        );
    }

    // Save calibration
    let calibration_data = CalibrationData {
        num_samples: train_data.len().min(50),
        layers,
        data_source: format!("ES.FUT ({})", dbn_file.display()),
        generated_at: chrono::Utc::now().to_rfc3339(),
    };

    let output_path = PathBuf::from("ml/checkpoints/tft_int8_calibration.json");
    if let Some(parent) = output_path.parent() {
        std::fs::create_dir_all(parent)?;
    }

    let json_string = serde_json::to_string_pretty(&calibration_data)?;
    std::fs::write(&output_path, json_string)?;

    let file_size = std::fs::metadata(&output_path)?.len();
    info!(
        "✅ Saved calibration to: {} ({} bytes)",
        output_path.display(),
        file_size
    );

    println!();
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!("  Calibration Complete!");
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!("  Output: {}", output_path.display());
    println!("  File size: {} bytes", file_size);
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

    Ok(())
}