╔═══════════════════════════════════════════════════════════════════════════╗ ║ TFT TRAINER SPLIT - VISUAL GUIDE ║ ╚═══════════════════════════════════════════════════════════════════════════╝ ┌───────────────────────────────────────────────────────────────────────────┐ │ ORIGINAL FILE: ml/src/trainers/tft.rs (2,915 lines) │ └───────────────────────────────────────────────────────────────────────────┘ Lines │ Content │ Destination ─────────┼────────────────────────────────┼───────────────────────────── 1-33 │ Header & Imports │ trainer.rs (modified) 34-100 │ QAT Metrics (5 structs) │ types.rs 101-202 │ TFTModel Trait │ model.rs 203-277 │ TFTTrainer Struct │ trainer.rs 278-332 │ TrainingState │ types.rs (pub(crate)) 333-385 │ Progress Types │ types.rs 386-538 │ TFTTrainerConfig │ config.rs 539-2498 │ TFTTrainer impl (30 methods) │ trainer.rs 2499-2507│ ValidationMetrics │ types.rs (pub(crate)) 2508-2548│ TrainingMetrics │ types.rs 2550-2915│ Test Module │ tests.rs ┌───────────────────────────────────────────────────────────────────────────┐ │ NEW STRUCTURE: ml/src/trainers/tft/ (6 files) │ └───────────────────────────────────────────────────────────────────────────┘ mod.rs (~50 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ //! Module documentation │ │ mod config; │ │ mod model; │ │ mod trainer; │ │ mod types; │ │ #[cfg(test)] mod tests; │ │ │ │ pub use config::TFTTrainerConfig; │ │ pub use model::TFTModel; │ │ pub use trainer::TFTTrainer; │ │ pub use types::{QATMetrics, TrainingProgress, ...}; │ └─────────────────────────────────────────────────────────────────────────┘ config.rs (~350 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ use serde::{Deserialize, Serialize}; │ │ use crate::tft::{TFTConfig, training::TFTTrainingConfig}; │ │ │ │ pub struct TFTTrainerConfig { │ │ pub epochs: usize, │ │ pub learning_rate: f64, │ │ pub batch_size: usize, │ │ // ... ~40 fields total │ │ } │ │ │ │ impl Default for TFTTrainerConfig { ... } │ │ │ │ impl TFTTrainerConfig { │ │ pub fn to_model_config(&self) -> TFTConfig { ... } │ │ pub fn to_training_config(&self) -> TFTTrainingConfig { ... } │ │ } │ └─────────────────────────────────────────────────────────────────────────┘ types.rs (~250 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ use std::collections::HashMap; │ │ use serde::{Deserialize, Serialize}; │ │ │ │ // PUBLIC: QAT Metrics (5 structs, ~67 lines) │ │ pub struct QATMetrics { ... } │ │ pub struct ScaleStatistics { ... } │ │ pub struct ZeroPointStatistics { ... } │ │ pub struct ObserverRangeStatistics { ... } │ │ pub struct LayerQuantizationMetrics { ... } │ │ │ │ // PUBLIC: Progress Types (~53 lines) │ │ pub struct TrainingProgress { ... } │ │ pub struct ResourceUsage { ... } │ │ │ │ // PUBLIC: Result Types (~41 lines) │ │ pub struct TrainingMetrics { ... } │ │ │ │ // PRIVATE: Internal State (~64 lines) │ │ pub(crate) struct TrainingState { ... } │ │ pub(crate) struct ValidationMetrics { ... } │ └─────────────────────────────────────────────────────────────────────────┘ model.rs (~100 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ use std::sync::Arc; │ │ use candle_core::{Device, Tensor}; │ │ use candle_nn::VarMap; │ │ use crate::tft::{TFTConfig, TemporalFusionTransformer}; │ │ use crate::MLError; │ │ │ │ pub trait TFTModel: Send + Sync { │ │ fn forward(...) -> Result; │ │ fn get_device(&self) -> &Device; │ │ fn get_config(&self) -> &TFTConfig; │ │ fn get_varmap(&self) -> Arc; │ │ fn clear_cache(&mut self); │ │ } │ │ │ │ impl TFTModel for TemporalFusionTransformer { ... } │ │ │ │ /* Future QAT implementation */ │ └─────────────────────────────────────────────────────────────────────────┘ trainer.rs (~2,000 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ use std::collections::HashMap; │ │ use std::sync::Arc; │ │ // ... external imports ... │ │ │ │ use super::config::TFTTrainerConfig; │ │ use super::model::TFTModel; │ │ use super::types::{QATMetrics, TrainingProgress, ...}; │ │ │ │ pub struct TFTTrainer { │ │ model_config: TFTConfig, │ │ training_config: TFTTrainingConfig, │ │ model: Box, │ │ optimizer: Option, │ │ // ... ~20 more fields │ │ } │ │ │ │ impl TFTTrainer { │ │ // Constructor │ │ pub fn new(...) -> MLResult { ... } │ │ │ │ // Public API (3 methods) │ │ pub fn set_progress_callback(...) { ... } │ │ pub async fn train(...) -> MLResult { ... } │ │ │ │ // Training Methods (6 methods) │ │ async fn train_epoch(...) -> MLResult { ... } │ │ async fn validate_epoch(...) -> MLResult { ... } │ │ fn batch_to_tensors(...) { ... } │ │ fn compute_quantile_loss(...) { ... } │ │ fn compute_rmse(...) { ... } │ │ fn extract_attention_entropy(...) { ... } │ │ │ │ // QAT Methods (5 methods) │ │ async fn run_qat_calibration(...) { ... } │ │ async fn qat_to_quantized_checkpoint(...) { ... } │ │ fn export_qat_metrics(...) { ... } │ │ fn apply_qat_lr_schedule(...) { ... } │ │ async fn quantize_and_save_int8_checkpoint(...) { ... } │ │ │ │ // Helper Methods (15+ methods) │ │ fn initialize_optimizer(...) { ... } │ │ fn check_early_stopping(...) { ... } │ │ async fn save_checkpoint(...) { ... } │ │ async fn send_progress_update(...) { ... } │ │ fn get_resource_usage(...) { ... } │ │ // ... more helpers ... │ │ } │ └─────────────────────────────────────────────────────────────────────────┘ tests.rs (~165 lines) ═══════════════════════════════════════════════════════════════════════════ ┌─────────────────────────────────────────────────────────────────────────┐ │ use super::*; │ │ use crate::checkpoint::FileSystemStorage; │ │ use std::path::PathBuf; │ │ use tempfile::TempDir; │ │ │ │ #[tokio::test] │ │ async fn test_tft_trainer_creation() { ... } │ │ │ │ #[tokio::test] │ │ async fn test_training_config_conversion() { ... } │ │ │ │ #[tokio::test] │ │ async fn test_checkpoint_save_load() { ... } │ │ │ │ #[tokio::test] │ │ async fn test_oom_retry_exponential_backoff() { ... } │ │ │ │ #[tokio::test] │ │ async fn test_oom_retry_minimum_batch_size() { ... } │ └─────────────────────────────────────────────────────────────────────────┘ ┌───────────────────────────────────────────────────────────────────────────┐ │ DEPENDENCY FLOW │ └───────────────────────────────────────────────────────────────────────────┘ ┌─────────────┐ │ mod.rs │ (orchestrates all) └──────┬──────┘ │ ┌────────────────┼────────────────┐ │ │ │ ▼ ▼ ▼ ┌──────────┐ ┌──────────┐ ┌──────────┐ │config.rs │ │ types.rs │ │ model.rs │ └──────────┘ └──────────┘ └──────────┘ │ │ │ └────────────────┼────────────────┘ │ ▼ ┌─────────────┐ │ trainer.rs │ (uses all above) └──────┬──────┘ │ ▼ ┌─────────────┐ │ tests.rs │ (tests trainer) └─────────────┘ ┌───────────────────────────────────────────────────────────────────────────┐ │ EXTRACTION CHECKLIST │ └───────────────────────────────────────────────────────────────────────────┘ Step 1: Create types.rs ✓ Extract lines 34-100 (QAT metrics) ✓ Extract lines 278-332 (TrainingState - make pub(crate)) ✓ Extract lines 333-385 (Progress types) ✓ Extract lines 2,499-2,507 (ValidationMetrics - make pub(crate)) ✓ Extract lines 2,508-2,548 (TrainingMetrics) ✓ Add module header and imports Step 2: Create model.rs ✓ Extract lines 101-202 (TFTModel trait + impls) ✓ Add module header and imports Step 3: Create config.rs ✓ Extract lines 386-538 (TFTTrainerConfig) ✓ Add module header and imports Step 4: Create trainer.rs ✓ Copy header from lines 1-33 (modified) ✓ Extract lines 203-277 (TFTTrainer struct) ✓ Extract lines 539-2,498 (impl blocks) ✓ Update imports to use super::{config, model, types} Step 5: Create tests.rs ✓ Extract lines 2,550-2,915 (test module) ✓ Remove #[cfg(test)] attribute ✓ Add use super::*; Step 6: Create mod.rs ✓ Write new file with module declarations ✓ Add re-exports for public API Step 7: Verify ✓ cargo check --package ml ✓ cargo test --package ml trainers::tft ✓ cargo clippy --package ml ╔═══════════════════════════════════════════════════════════════════════════╗ ║ READY TO PROCEED ║ ║ ║ ║ All documentation complete. Implementation agent can begin extraction. ║ ╚═══════════════════════════════════════════════════════════════════════════╝