candle_coreml/

model_config.rs

//! Model configuration system for candle-coreml
//!
//! This module provides a configuration system that replaces hardcoded model shapes
//! with discoverable, flexible configurations. It supports loading configurations from
//! JSON files generated by the shape discovery tool, as well as built-in configurations
//! for known models.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// Complete model configuration including shapes, components, and naming patterns
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelConfig {
    pub model_info: ModelInfo,
    pub shapes: ShapeConfig,
    pub components: HashMap<String, ComponentConfig>,
    pub naming: NamingConfig,
    /// Execution mode for FFN: "unified" (single component/function) or "split" (separate prefill/infer components)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_execution: Option<String>,
}

/// Model metadata and identification
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ModelInfo {
    #[serde(default)]
    pub model_id: Option<String>,
    pub path: Option<String>,
    pub model_type: String,
    pub discovered_at: Option<String>,
}

/// Overall model shape parameters
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ShapeConfig {
    pub batch_size: usize,
    pub context_length: usize,
    pub hidden_size: usize,
    pub vocab_size: usize,
}

/// Configuration for a single model component (embeddings, FFN, LM head, etc.)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ComponentConfig {
    pub file_path: Option<String>,
    pub inputs: HashMap<String, TensorConfig>,
    pub outputs: HashMap<String, TensorConfig>,
    pub functions: Vec<String>,
    /// Optional deterministic input order; if absent, caller must provide correct order.
    #[serde(default)]
    pub input_order: Option<Vec<String>>,
}

/// Configuration for a tensor (shape and data type)
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct TensorConfig {
    pub name: String,
    pub shape: Vec<usize>,
    pub data_type: String,
}

/// Model file naming patterns for component discovery
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct NamingConfig {
    // Deprecated: patterns removed; explicit file paths are required per component
    #[serde(skip_serializing_if = "Option::is_none")]
    pub embeddings_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_prefill_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ffn_infer_pattern: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub lm_head_pattern: Option<String>,
}

impl ModelConfig {
    /// Load configuration from a JSON file
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let path = path.as_ref();
        let content = std::fs::read_to_string(path)
            .with_context(|| format!("Failed to read config file: {}", path.display()))?;

        let config: ModelConfig = serde_json::from_str(&content)
            .with_context(|| format!("Failed to parse config file: {}", path.display()))?;

        Ok(config)
    }

    /// Save configuration to a JSON file
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let path = path.as_ref();
        let content =
            serde_json::to_string_pretty(self).context("Failed to serialize configuration")?;

        std::fs::write(path, content)
            .with_context(|| format!("Failed to write config file: {}", path.display()))?;

        Ok(())
    }

    /// Get built-in configuration for known model IDs
    pub fn get_builtin_config(model_id: &str) -> Option<Self> {
        crate::builtin_configs::get_builtin_config(model_id)
    }

    /// Create a default configuration with standard Qwen shapes
    pub fn default_qwen() -> Self {
        Self {
            model_info: ModelInfo {
                model_id: None,
                path: None,
                model_type: "qwen".to_string(),
                discovered_at: None,
            },
            shapes: ShapeConfig {
                batch_size: 1,
                context_length: 512,
                hidden_size: 1024,
                vocab_size: 151936,
            },
            components: HashMap::new(),
            naming: NamingConfig {
                embeddings_pattern: None,
                ffn_prefill_pattern: None,
                ffn_infer_pattern: None,
                lm_head_pattern: None,
            },
            ffn_execution: None,
        }
    }

    /// Get the shape configuration for a specific component and tensor
    pub fn get_tensor_shape(
        &self,
        component: &str,
        tensor_name: &str,
        is_input: bool,
    ) -> Option<&Vec<usize>> {
        let component_config = self.components.get(component)?;

        let tensor_map = if is_input {
            &component_config.inputs
        } else {
            &component_config.outputs
        };

        tensor_map.get(tensor_name).map(|tensor| &tensor.shape)
    }

    /// Get the expected input shape for embeddings
    pub fn embeddings_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("embeddings", "input_ids", true)
    }

    /// Get the expected output shape for embeddings
    pub fn embeddings_output_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("embeddings", "hidden_states", false)
    }

    /// Get the expected input shape for FFN prefill
    pub fn ffn_prefill_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("ffn_prefill", "hidden_states", true)
    }

    /// Get the expected input shape for LM head
    pub fn lm_head_input_shape(&self) -> Option<&Vec<usize>> {
        self.get_tensor_shape("lm_head", "hidden_states", true)
    }

    /// Check if this model supports multi-part logits output
    pub fn has_multipart_logits(&self) -> bool {
        if let Some(lm_head) = self.components.get("lm_head") {
            // Check if there are multiple logits outputs (logits1, logits2, etc.)
            let logits_outputs: Vec<_> = lm_head
                .outputs
                .keys()
                .filter(|name| name.starts_with("logits") && name.len() > 6) // "logits" + number
                .collect();
            return logits_outputs.len() > 1;
        }
        false
    }

    /// Get the number of logits parts for multi-part output
    pub fn logits_part_count(&self) -> usize {
        if let Some(lm_head) = self.components.get("lm_head") {
            let logits_outputs: Vec<_> = lm_head
                .outputs
                .keys()
                .filter(|name| name.starts_with("logits"))
                .collect();
            if logits_outputs.is_empty() {
                1 // Single logits output
            } else {
                logits_outputs.len()
            }
        } else {
            1
        }
    }

    /// Select the primary logits output name for the LM head.
    /// Preference order: "logits1" (multipart), then "logits", otherwise the first available key.
    pub fn lm_head_primary_output_name(&self) -> Option<String> {
        let lm_head = self.components.get("lm_head")?;

        // Prefer explicit multipart naming starting with logits1
        if lm_head.outputs.contains_key("logits1") {
            return Some("logits1".to_string());
        }

        // Common single output name
        if lm_head.outputs.contains_key("logits") {
            return Some("logits".to_string());
        }

        // Fallback to the first key if any
        lm_head.outputs.keys().next().map(|k| k.to_string())
    }

    /// Validate the configuration for consistency
    pub fn validate(&self) -> Result<()> {
        // Check that required components exist
        let required_components = ["embeddings", "lm_head"];
        for component in required_components {
            if !self.components.contains_key(component) {
                return Err(anyhow::anyhow!("Missing required component: {}", component));
            }
        }

        // Check shape consistency
        if self.shapes.batch_size == 0 {
            return Err(anyhow::anyhow!("batch_size must be greater than 0"));
        }

        if self.shapes.context_length == 0 {
            return Err(anyhow::anyhow!("context_length must be greater than 0"));
        }

        if self.shapes.hidden_size == 0 {
            return Err(anyhow::anyhow!("hidden_size must be greater than 0"));
        }

        if self.shapes.vocab_size == 0 {
            return Err(anyhow::anyhow!("vocab_size must be greater than 0"));
        }

        // Validate tensor shapes make sense
        for (component_name, component) in &self.components {
            for (tensor_name, tensor) in &component.inputs {
                if tensor.shape.is_empty() {
                    return Err(anyhow::anyhow!(
                        "Empty shape for {}.inputs.{}",
                        component_name,
                        tensor_name
                    ));
                }
            }
            for (tensor_name, tensor) in &component.outputs {
                if tensor.shape.is_empty() {
                    return Err(anyhow::anyhow!(
                        "Empty shape for {}.outputs.{}",
                        component_name,
                        tensor_name
                    ));
                }
            }
        }

        Ok(())
    }

    /// Validate internal wiring between components for basic shape compatibility.
    /// Examples:
    ///  - embeddings.outputs.hidden_states == ffn_prefill.inputs.hidden_states
    ///  - ffn_infer.outputs.output_hidden_states == lm_head.inputs.hidden_states (when ffn_infer exists)
    pub fn validate_internal_wiring(&self) -> Result<()> {
        // Embeddings -> FFN prefill hidden_states flow
        if let (Some(emb_out), Some(ffn_in_hidden)) = (
            self.get_tensor_shape("embeddings", "hidden_states", false),
            self.get_tensor_shape("ffn_prefill", "hidden_states", true),
        ) {
            if emb_out != ffn_in_hidden {
                return Err(anyhow::anyhow!(
                    "Shape mismatch: embeddings.hidden_states {:?} != ffn_prefill.hidden_states {:?}",
                    emb_out, ffn_in_hidden
                ));
            }
        }

        // FFN infer -> LM head hidden_states flow
        if self.components.contains_key("ffn_infer") {
            if let (Some(ffn_out), Some(lm_in)) = (
                self.get_tensor_shape("ffn_infer", "output_hidden_states", false),
                self.get_tensor_shape("lm_head", "hidden_states", true),
            ) {
                if ffn_out != lm_in {
                    return Err(anyhow::anyhow!(
                        "Shape mismatch: ffn_infer.output_hidden_states {:?} != lm_head.hidden_states {:?}",
                        ffn_out, lm_in
                    ));
                }
            }
        } else {
            // If there's no separate ffn_infer, check ffn_prefill output shape matches lm_head (single-token path)
            if let (Some(ffn_out), Some(lm_in)) = (
                self.get_tensor_shape("ffn_prefill", "output_hidden_states", false),
                self.get_tensor_shape("lm_head", "hidden_states", true),
            ) {
                if ffn_out != lm_in {
                    return Err(anyhow::anyhow!(
                        "Shape mismatch: ffn_prefill.output_hidden_states {:?} != lm_head.hidden_states {:?}",
                        ffn_out, lm_in
                    ));
                }
            }
        }

        Ok(())
    }

    /// Determine if FFN execution should be treated as split (separate infer component)
    pub fn ffn_is_split(&self) -> bool {
        if let Some(mode) = self.ffn_execution.as_deref() {
            return mode == "split";
        }
        if let (Some(prefill), Some(infer)) = (
            self.components.get("ffn_prefill"),
            self.components.get("ffn_infer"),
        ) {
            match (&prefill.file_path, &infer.file_path) {
                (Some(p), Some(i)) => p != i, // different files => split
                _ => false,
            }
        } else {
            false
        }
    }

    /// Detect if prefill should run in single-token sequential mode based on configured shapes
    pub fn prefill_is_single_token(&self) -> bool {
        if let Some(prefill) = self.components.get("ffn_prefill") {
            if let Some(hs) = prefill.inputs.get("hidden_states") {
                return hs.shape.len() == 3 && hs.shape.get(1) == Some(&1);
            }
        }
        false
    }
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self::default_qwen()
    }
}

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

    fn create_test_config() -> ModelConfig {
        let mut components = HashMap::new();

        // Embeddings component
        let mut embeddings_inputs = HashMap::new();
        embeddings_inputs.insert(
            "input_ids".to_string(),
            TensorConfig {
                name: "input_ids".to_string(),
                shape: vec![1, 64],
                data_type: "INT32".to_string(),
            },
        );

        let mut embeddings_outputs = HashMap::new();
        embeddings_outputs.insert(
            "hidden_states".to_string(),
            TensorConfig {
                name: "hidden_states".to_string(),
                shape: vec![1, 64, 1024],
                data_type: "FLOAT16".to_string(),
            },
        );

        components.insert(
            "embeddings".to_string(),
            ComponentConfig {
                file_path: None,
                inputs: embeddings_inputs,
                outputs: embeddings_outputs,
                functions: vec![],
                input_order: None,
            },
        );

        // LM Head component
        let mut lm_head_inputs = HashMap::new();
        lm_head_inputs.insert(
            "hidden_states".to_string(),
            TensorConfig {
                name: "hidden_states".to_string(),
                shape: vec![1, 1, 1024],
                data_type: "FLOAT16".to_string(),
            },
        );

        let mut lm_head_outputs = HashMap::new();
        lm_head_outputs.insert(
            "logits".to_string(),
            TensorConfig {
                name: "logits".to_string(),
                shape: vec![1, 1, 151936],
                data_type: "FLOAT32".to_string(),
            },
        );

        components.insert(
            "lm_head".to_string(),
            ComponentConfig {
                file_path: None,
                inputs: lm_head_inputs,
                outputs: lm_head_outputs,
                functions: vec![],
                input_order: None,
            },
        );

        ModelConfig {
            model_info: ModelInfo {
                model_id: Some("test/model".to_string()),
                path: Some("/test/path".to_string()),
                model_type: "qwen".to_string(),
                discovered_at: Some("2025-08-07T00:00:00".to_string()),
            },
            shapes: ShapeConfig {
                batch_size: 1,
                context_length: 512,
                hidden_size: 1024,
                vocab_size: 151936,
            },
            components,
            naming: NamingConfig {
                embeddings_pattern: None,
                ffn_prefill_pattern: None,
                ffn_infer_pattern: None,
                lm_head_pattern: None,
            },
            ffn_execution: Some("unified".to_string()),
        }
    }

    #[test]
    fn test_config_serialization() {
        let config = create_test_config();

        // Test JSON serialization
        let json = serde_json::to_string_pretty(&config).unwrap();
        assert!(json.contains("test/model"));
        assert!(json.contains("batch_size"));
        assert!(json.contains("embeddings"));

        // Test deserialization
        let parsed: ModelConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(parsed.model_info.model_id, config.model_info.model_id);
        assert_eq!(parsed.shapes.batch_size, config.shapes.batch_size);
        assert_eq!(parsed.components.len(), config.components.len());
    }

    #[test]
    fn test_config_file_io() {
        let config = create_test_config();
        let temp_file = NamedTempFile::new().unwrap();

        // Save configuration
        config.save_to_file(temp_file.path()).unwrap();

        // Load configuration
        let loaded = ModelConfig::load_from_file(temp_file.path()).unwrap();
        assert_eq!(loaded.model_info.model_id, config.model_info.model_id);
        assert_eq!(loaded.shapes.hidden_size, config.shapes.hidden_size);
    }

    #[test]
    fn test_shape_accessors() {
        let config = create_test_config();

        // Test embeddings shapes
        let embeddings_input = config.embeddings_input_shape().unwrap();
        assert_eq!(embeddings_input, &vec![1, 64]);

        let embeddings_output = config.embeddings_output_shape().unwrap();
        assert_eq!(embeddings_output, &vec![1, 64, 1024]);

        let lm_head_input = config.lm_head_input_shape().unwrap();
        assert_eq!(lm_head_input, &vec![1, 1, 1024]);
    }

    #[test]
    fn test_multipart_logits_detection() {
        let config = create_test_config();
        assert!(!config.has_multipart_logits()); // Single logits output

        // Create config with multipart logits
        let mut config_multipart = config;
        let lm_head = config_multipart.components.get_mut("lm_head").unwrap();
        lm_head.outputs.clear();
        lm_head.outputs.insert(
            "logits1".to_string(),
            TensorConfig {
                name: "logits1".to_string(),
                shape: vec![1, 1, 9480],
                data_type: "FLOAT32".to_string(),
            },
        );
        lm_head.outputs.insert(
            "logits2".to_string(),
            TensorConfig {
                name: "logits2".to_string(),
                shape: vec![1, 1, 9479],
                data_type: "FLOAT32".to_string(),
            },
        );

        assert!(config_multipart.has_multipart_logits());
        assert_eq!(config_multipart.logits_part_count(), 2);
    }

    #[test]
    fn test_config_validation() {
        let config = create_test_config();
        assert!(config.validate().is_ok());

        // Internal wiring should be consistent in this synthetic setup
        assert!(config.validate_internal_wiring().is_ok());

        // Test missing component
        let mut invalid_config = config.clone();
        invalid_config.components.remove("embeddings");
        assert!(invalid_config.validate().is_err());

        // Test invalid shapes
        let mut invalid_shapes = config;
        invalid_shapes.shapes.batch_size = 0;
        assert!(invalid_shapes.validate().is_err());
    }

    #[test]
    fn test_default_config() {
        let config = ModelConfig::default();
        assert_eq!(config.model_info.model_type, "qwen");
        assert_eq!(config.shapes.batch_size, 1);
        assert_eq!(config.shapes.context_length, 512);
        assert_eq!(config.shapes.hidden_size, 1024);
        assert_eq!(config.shapes.vocab_size, 151936);
    }
}