ohms_adaptq/

streaming_loader.rs

use crate::{Result, WeightMatrix};
use reqwest::Client;
use std::io::{Read, Seek, SeekFrom, Cursor};
use std::collections::HashMap;
use serde_json::Value;
use memmap2::MmapOptions;
use std::fs::File;
use tempfile::NamedTempFile;

/// Streaming model loader that processes models directly from online sources
/// without requiring full download to disk
#[derive(Debug)]
pub struct StreamingModelLoader {
    client: Client,
    chunk_size: usize,
    temp_threshold: usize, // Only use temp files for models larger than this
}

impl StreamingModelLoader {
    pub fn new() -> Self {
        Self {
            client: Client::builder()
                .timeout(std::time::Duration::from_secs(300))
                .build()
                .expect("Failed to create HTTP client"),
            chunk_size: 64 * 1024 * 1024, // 64MB chunks
            temp_threshold: 512 * 1024 * 1024, // 512MB threshold
        }
    }

    /// Load model weights with streaming compression
    pub async fn load_and_compress_streaming(
        &self, 
        source_url: &str,
        novaq_engine: &mut crate::NOVAQEngine,
        progress_callback: impl Fn(f32, &str),
    ) -> Result<crate::NOVAQModel> {
        progress_callback(0.0, "Fetching model metadata...");
        
        // First, get model metadata without downloading the full model
        let metadata = self.fetch_model_metadata(source_url).await?;
        
        match metadata.format {
            ModelFormat::SafeTensors => {
                self.stream_compress_safetensors(source_url, novaq_engine, progress_callback).await
            }
            ModelFormat::PyTorch => {
                self.stream_compress_pytorch(source_url, novaq_engine, progress_callback).await
            }
            ModelFormat::GGUF => {
                self.stream_compress_gguf(source_url, novaq_engine, progress_callback).await
            }
            ModelFormat::HuggingFace => {
                self.stream_compress_huggingface(source_url, novaq_engine, progress_callback).await
            }
        }
    }

    /// Fetch only model metadata without downloading the full model
    async fn fetch_model_metadata(&self, url: &str) -> Result<ModelMetadata> {
        // For SafeTensors, we only need to read the header
        if url.ends_with(".safetensors") {
            return self.fetch_safetensors_metadata(url).await;
        }
        
        // For HuggingFace repos, check config.json and model files
        if url.contains("huggingface.co") || url.contains("hf:") {
            return self.fetch_huggingface_metadata(url).await;
        }
        
        // For other formats, we might need to download a small portion
        self.fetch_generic_metadata(url).await
    }

    /// Stream compress SafeTensors format
    async fn stream_compress_safetensors(
        &self,
        url: &str,
        novaq_engine: &mut crate::NOVAQEngine,
        progress_callback: impl Fn(f32, &str),
    ) -> Result<crate::NOVAQModel> {
        progress_callback(0.1, "Reading SafeTensors header...");
        
        // Read just the header first (first 8 bytes + header size)
        let mut response = self.client.get(url).send().await?;
        let total_size = response.content_length().unwrap_or(0);
        
        let mut header_buffer = Vec::new();
        let mut bytes_read = 0;
        
        // Read header size (first 8 bytes)
        let mut size_bytes = [0u8; 8];
        response.chunk().await?;
        // Read the actual header size and parse it
        
        progress_callback(0.2, "Parsing model structure...");
        
        // Parse SafeTensors header to get tensor information
        let header: HashMap<String, Value> = serde_json::from_slice(&header_buffer)?;
        
        let mut weight_matrices = Vec::new();
        let mut tensor_count = 0;
        let total_tensors = header.len();
        
        progress_callback(0.3, "Starting streaming compression...");
        
        // Process each tensor in chunks without storing the full model
        for (tensor_name, tensor_info) in header.iter() {
            if tensor_name == "__metadata__" {
                continue;
            }
            
            tensor_count += 1;
            let progress = 0.3 + (tensor_count as f32 / total_tensors as f32) * 0.6;
            progress_callback(progress, &format!("Compressing tensor: {}", tensor_name));
            
            // Extract tensor data in chunks
            let tensor_data = self.extract_tensor_streaming(url, tensor_info).await?;
            
            // Create weight matrix
            let shape = tensor_info["shape"].as_array()
                .ok_or("Invalid tensor shape")?
                .iter()
                .map(|v| v.as_u64().unwrap() as usize)
                .collect();
            
            let weight_matrix = WeightMatrix::new(tensor_data, shape, tensor_name.clone());
            weight_matrices.push(weight_matrix);
        }
        
        progress_callback(0.9, "Finalizing NOVAQ compression...");
        
        // Process all weight matrices with NOVAQ
        let novaq_model = novaq_engine.quantize_model(weight_matrices)?;
        
        progress_callback(1.0, "Streaming compression complete!");
        
        Ok(novaq_model)
    }

    /// Extract tensor data in streaming fashion
    async fn extract_tensor_streaming(&self, url: &str, tensor_info: &Value) -> Result<Vec<f32>> {
        let data_offsets = tensor_info["data_offsets"].as_array()
            .ok_or("Invalid tensor data offsets")?;
        let start_offset = data_offsets[0].as_u64().unwrap();
        let end_offset = data_offsets[1].as_u64().unwrap();
        let tensor_size = (end_offset - start_offset) as usize;
        
        // Use Range request to get only the tensor data we need
        let response = self.client
            .get(url)
            .header("Range", format!("bytes={}-{}", start_offset, end_offset - 1))
            .send()
            .await?;
        
        let tensor_bytes = response.bytes().await?;
        
        // Convert bytes to f32 values (assuming fp32)
        let float_count = tensor_size / 4;
        let mut tensor_data = Vec::with_capacity(float_count);
        
        for chunk in tensor_bytes.chunks_exact(4) {
            let float_val = f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
            tensor_data.push(float_val);
        }
        
        Ok(tensor_data)
    }

    /// Stream compress PyTorch format
    async fn stream_compress_pytorch(
        &self,
        url: &str,
        novaq_engine: &mut crate::NOVAQEngine,
        progress_callback: impl Fn(f32, &str),
    ) -> Result<crate::NOVAQModel> {
        progress_callback(0.0, "PyTorch streaming not yet implemented");
        // TODO: Implement PyTorch streaming compression
        Err("PyTorch streaming compression not implemented".into())
    }

    /// Stream compress GGUF format  
    async fn stream_compress_gguf(
        &self,
        url: &str,
        novaq_engine: &mut crate::NOVAQEngine,
        progress_callback: impl Fn(f32, &str),
    ) -> Result<crate::NOVAQModel> {
        progress_callback(0.0, "GGUF streaming not yet implemented");
        // TODO: Implement GGUF streaming compression
        Err("GGUF streaming compression not implemented".into())
    }

    /// Stream compress HuggingFace repository
    async fn stream_compress_huggingface(
        &self,
        url: &str,
        novaq_engine: &mut crate::NOVAQEngine,
        progress_callback: impl Fn(f32, &str),
    ) -> Result<crate::NOVAQModel> {
        progress_callback(0.0, "Detecting HuggingFace model files...");
        
        // Find the main model file in the HF repo
        let model_files = self.list_huggingface_files(url).await?;
        
        // Prioritize SafeTensors, then PyTorch
        let model_file = model_files.iter()
            .find(|f| f.ends_with(".safetensors"))
            .or_else(|| model_files.iter().find(|f| f.ends_with(".bin")))
            .ok_or("No supported model files found in HuggingFace repo")?;
        
        let full_url = if url.starts_with("hf:") {
            // Convert hf:repo/model to actual HF URL
            let repo = url.strip_prefix("hf:").unwrap();
            format!("https://huggingface.co/{}/resolve/main/{}", repo, model_file)
        } else {
            format!("{}/resolve/main/{}", url, model_file)
        };
        
        // Stream compress the actual model file
        if model_file.ends_with(".safetensors") {
            self.stream_compress_safetensors(&full_url, novaq_engine, progress_callback).await
        } else {
            self.stream_compress_pytorch(&full_url, novaq_engine, progress_callback).await
        }
    }

    /// Fetch SafeTensors metadata without downloading full file
    async fn fetch_safetensors_metadata(&self, url: &str) -> Result<ModelMetadata> {
        // Read just the first 1KB to get header size
        let response = self.client
            .get(url)
            .header("Range", "bytes=0-1023")
            .send()
            .await?;
        
        let initial_bytes = response.bytes().await?;
        
        // Parse header size from first 8 bytes
        let header_size = u64::from_le_bytes([
            initial_bytes[0], initial_bytes[1], initial_bytes[2], initial_bytes[3],
            initial_bytes[4], initial_bytes[5], initial_bytes[6], initial_bytes[7],
        ]) as usize;
        
        Ok(ModelMetadata {
            format: ModelFormat::SafeTensors,
            estimated_size: 0, // Will be determined during streaming
            tensor_count: 0,   // Will be determined from header
        })
    }

    /// Fetch HuggingFace metadata
    async fn fetch_huggingface_metadata(&self, url: &str) -> Result<ModelMetadata> {
        Ok(ModelMetadata {
            format: ModelFormat::HuggingFace,
            estimated_size: 0,
            tensor_count: 0,
        })
    }

    /// Fetch generic metadata for unknown formats
    async fn fetch_generic_metadata(&self, url: &str) -> Result<ModelMetadata> {
        // Try to determine format from URL extension
        let format = if url.ends_with(".safetensors") {
            ModelFormat::SafeTensors
        } else if url.ends_with(".bin") || url.ends_with(".pt") || url.ends_with(".pth") {
            ModelFormat::PyTorch
        } else if url.ends_with(".gguf") {
            ModelFormat::GGUF
        } else {
            ModelFormat::SafeTensors // Default assumption
        };
        
        Ok(ModelMetadata {
            format,
            estimated_size: 0,
            tensor_count: 0,
        })
    }

    /// List files in HuggingFace repository
    async fn list_huggingface_files(&self, url: &str) -> Result<Vec<String>> {
        // Simple implementation - in practice, would use HF API
        Ok(vec!["model.safetensors".to_string(), "pytorch_model.bin".to_string()])
    }
}

#[derive(Debug)]
pub struct ModelMetadata {
    pub format: ModelFormat,
    pub estimated_size: u64,
    pub tensor_count: usize,
}

#[derive(Debug)]
pub enum ModelFormat {
    SafeTensors,
    PyTorch,
    GGUF,
    HuggingFace,
}

impl Default for StreamingModelLoader {
    fn default() -> Self {
        Self::new()
    }
}