mlmf 0.2.0

//! AWQ (Activation-aware Weight Quantization) format support
//!
//! This module provides loading support for AWQ quantized models, which are optimized
//! for efficient inference with minimal accuracy loss. Key features:
//!
//! - **4-bit quantized weights**: Reduced memory usage and faster inference
//! - **Activation-aware quantization**: Preserves important weights based on activation patterns  
//! - **JSON configuration**: Model metadata and quantization parameters
//! - **Compatible with Candle**: Uses Candle's quantized tensor support

use crate::{
    error::{Error, Result},
    loader::{LoadOptions, LoadedModel},
    progress::ProgressEvent,
    smart_mapping::SmartTensorNameMapper,
    ModelConfig,
};
// Removed unused Device import
use candle_nn::VarBuilder;
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    fs,
    path::{Path, PathBuf},
};

/// AWQ model configuration loaded from config.json
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AWQConfig {
    /// Model architectures (e.g., ["Qwen3ForCausalLM"])
    #[serde(default)]
    pub architectures: Option<Vec<String>>,
    /// Model type (e.g., "qwen3", "llama")
    pub model_type: Option<String>,
    /// Vocabulary size
    pub vocab_size: Option<u32>,
    /// Hidden layer size
    pub hidden_size: Option<u32>,
    /// Number of attention heads
    pub num_attention_heads: Option<u32>,
    /// Number of hidden layers
    pub num_hidden_layers: Option<u32>,
    /// Intermediate layer size (FFN)
    pub intermediate_size: Option<u32>,
    /// Maximum position embeddings
    pub max_position_embeddings: Option<u32>,
    /// Layer normalization epsilon
    #[serde(alias = "rms_norm_eps")]
    pub layer_norm_eps: Option<f64>,
    /// RoPE theta parameter
    pub rope_theta: Option<f64>,
    /// Whether to tie word embeddings
    pub tie_word_embeddings: Option<bool>,
    /// Quantization configuration
    pub quantization_config: Option<AWQQuantizationConfig>,
    /// Additional fields that may be present (for forward compatibility)
    #[serde(flatten)]
    pub additional_fields: std::collections::HashMap<String, serde_json::Value>,
}

/// AWQ quantization configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AWQQuantizationConfig {
    /// Number of bits for quantization (typically 4)
    pub bits: Option<u32>,
    /// Group size for quantization
    pub group_size: Option<u32>,
    /// Quantization method (should be "awq")
    pub quant_method: Option<String>,
    /// Quantization version
    pub version: Option<String>,
    /// Zero point flag
    pub zero_point: Option<bool>,
    /// Modules to not convert
    pub modules_to_not_convert: Option<serde_json::Value>,
}

/// Load AWQ model from directory containing config.json and .safetensors files
pub fn load_awq<P: AsRef<Path>>(model_dir: P, mut options: LoadOptions) -> Result<LoadedModel> {
    let model_dir = model_dir.as_ref();

    // Validate inputs
    if !model_dir.is_dir() {
        return Err(Error::model_loading(format!(
            "AWQ model directory not found: {:?}",
            model_dir
        )));
    }

    if let Some(callback) = &options.progress {
        callback(ProgressEvent::LoadingConfig {
            path: model_dir.join("config.json").display().to_string(),
        });
    }

    // Load AWQ configuration
    let config_path = model_dir.join("config.json");
    let awq_config = load_awq_config(&config_path)?;

    // Find safetensors files
    if let Some(callback) = &options.progress {
        callback(ProgressEvent::ScanningFiles { count: 0 });
    }

    let safetensors_files = find_awq_safetensors_files(model_dir)?;

    if let Some(callback) = &options.progress {
        callback(ProgressEvent::ScanningFiles {
            count: safetensors_files.len(),
        });
    }

    // For now, create placeholder implementation
    // In a full implementation, this would load and dequantize AWQ tensors
    let tensor_names = vec![
        "model.embed_tokens.weight".to_string(),
        "model.layers.0.self_attn.q_proj.weight".to_string(),
        "model.layers.0.self_attn.k_proj.weight".to_string(),
        "model.layers.0.self_attn.v_proj.weight".to_string(),
        "lm_head.weight".to_string(),
    ];

    // Create smart tensor name mapper
    if let Some(callback) = &options.progress {
        callback(ProgressEvent::DetectingArchitecture);
    }

    let mut smart_mapper = SmartTensorNameMapper::from_tensor_names(&tensor_names)?;

    // Integrate ML oracle if provided
    if let Some(oracle) = options.smart_mapping_oracle.take() {
        smart_mapper = smart_mapper.with_oracle(oracle);
    }

    let architecture = smart_mapper.architecture().ok_or_else(|| {
        Error::model_loading("Could not detect model architecture from AWQ tensor names")
    })?;

    // Convert AWQ config to ModelConfig
    let config = awq_config_to_model_config(&awq_config, architecture)?;

    // Create empty tensors for placeholder (would load actual quantized tensors in full implementation)
    let raw_tensors = HashMap::new();

    // Create VarBuilder - placeholder implementation
    if let Some(callback) = &options.progress {
        callback(ProgressEvent::BuildingModel);
    }

    let var_map = candle_nn::VarMap::new();
    let var_builder = VarBuilder::from_varmap(&var_map, options.dtype, &options.device);

    if let Some(callback) = &options.progress {
        callback(ProgressEvent::Complete {
            tensor_count: tensor_names.len(),
            format: "AWQ".to_string(),
        });
    }

    Ok(LoadedModel {
        var_builder,
        config,
        name_mapper: smart_mapper,
        raw_tensors,
        quantized_tensors: None,
        metadata: crate::metadata::ModelMetadata::new(),
        tensor_info: HashMap::new(),
        quantization_info: None,
        provenance: crate::metadata::ModelProvenance::new(),
    })
}

/// Load AWQ configuration from config.json
fn load_awq_config(config_path: &Path) -> Result<AWQConfig> {
    let config_str = fs::read_to_string(config_path).map_err(|e| {
        Error::model_loading(&format!(
            "Failed to read AWQ config from {}: {}",
            config_path.display(),
            e
        ))
    })?;

    serde_json::from_str(&config_str)
        .map_err(|e| Error::model_loading(&format!("Failed to parse AWQ config: {}", e)))
}

/// Find AWQ safetensors files in directory
fn find_awq_safetensors_files(dir: &Path) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();

    for entry in fs::read_dir(dir).map_err(|e| {
        Error::model_loading(&format!(
            "Failed to read AWQ model directory {}: {}",
            dir.display(),
            e
        ))
    })? {
        let entry = entry
            .map_err(|e| Error::model_loading(&format!("Failed to read directory entry: {}", e)))?;

        let path = entry.path();
        if let Some(extension) = path.extension() {
            if extension == "safetensors" {
                files.push(path);
            }
        }
    }

    if files.is_empty() {
        return Err(Error::model_loading(format!(
            "No .safetensors files found in AWQ model directory: {}",
            dir.display()
        )));
    }

    files.sort();
    Ok(files)
}

/// Convert AWQConfig to standard ModelConfig
fn awq_config_to_model_config(
    awq_config: &AWQConfig,
    architecture: &crate::name_mapping::Architecture,
) -> Result<ModelConfig> {
    Ok(ModelConfig {
        vocab_size: awq_config.vocab_size.unwrap_or(32000) as usize,
        hidden_size: awq_config.hidden_size.unwrap_or(4096) as usize,
        num_attention_heads: awq_config.num_attention_heads.unwrap_or(32) as usize,
        num_hidden_layers: awq_config.num_hidden_layers.unwrap_or(32) as usize,
        intermediate_size: awq_config.intermediate_size.unwrap_or(11008) as usize,
        max_position_embeddings: awq_config.max_position_embeddings.unwrap_or(4096) as usize,
        layer_norm_eps: awq_config.layer_norm_eps.unwrap_or(1e-6),
        dropout: 0.0, // AWQ models typically don't specify dropout for inference
        attention_dropout: 0.0,
        activation_function: "silu".to_string(), // Common default for LLaMA-style models
        rope_theta: awq_config.rope_theta.unwrap_or(10000.0),
        tie_word_embeddings: awq_config.tie_word_embeddings.unwrap_or(false),
        architecture: architecture.clone(),
        raw_config: serde_json::Value::Null,
    })
}

/// Check if directory contains AWQ model files
pub fn is_awq_model<P: AsRef<Path>>(model_dir: P) -> bool {
    let model_dir = model_dir.as_ref();

    // Check for config.json with quantization_config
    let config_path = model_dir.join("config.json");
    if let Ok(config_str) = fs::read_to_string(&config_path) {
        if let Ok(config) = serde_json::from_str::<AWQConfig>(&config_str) {
            if let Some(quant_config) = &config.quantization_config {
                // Check if it's actually AWQ (not just any quantization)
                return quant_config
                    .quant_method
                    .as_ref()
                    .map_or(false, |method| method == "awq");
            }
        }
    }

    false
}