xserv/crates/xserv-model/src/config.rs

use serde::Deserialize;
use std::path::Path;

#[derive(Debug, Clone, Deserialize)]
pub struct RopeScaling {
    pub rope_type: Option<String>,
    pub factor: Option<f64>,
    pub original_max_position_embeddings: Option<usize>,
    pub beta_fast: Option<f64>,
    pub beta_slow: Option<f64>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct ModelConfig {
    pub architectures: Option<Vec<String>>,
    pub model_type: Option<String>,

    // Modern HF naming
    #[serde(default)]
    pub hidden_size: Option<usize>,
    #[serde(default)]
    pub intermediate_size: Option<usize>,
    #[serde(default)]
    pub num_attention_heads: Option<usize>,
    #[serde(default)]
    pub num_key_value_heads: Option<usize>,
    #[serde(default)]
    pub num_hidden_layers: Option<usize>,
    pub vocab_size: usize,
    #[serde(default)]
    pub max_position_embeddings: Option<usize>,

    // GPT-2 naming
    #[serde(default)]
    pub n_embd: Option<usize>,
    #[serde(default)]
    pub n_head: Option<usize>,
    #[serde(default)]
    pub n_layer: Option<usize>,
    #[serde(default)]
    pub n_positions: Option<usize>,
    #[serde(default)]
    pub n_inner: Option<usize>,

    // Normalization
    #[serde(default)]
    pub layer_norm_eps: Option<f64>,
    #[serde(default)]
    pub layer_norm_epsilon: Option<f64>,
    #[serde(default)]
    pub rms_norm_eps: Option<f64>,

    // Other
    #[serde(default)]
    pub rope_theta: Option<f64>,
    #[serde(default)]
    pub tie_word_embeddings: Option<bool>,

    // MoE (gpt-oss)
    #[serde(default)]
    pub num_local_experts: Option<usize>,
    #[serde(default)]
    pub num_experts_per_tok: Option<usize>,
    #[serde(default)]
    pub layer_types: Option<Vec<String>>,
    #[serde(default)]
    pub sliding_window: Option<usize>,
    #[serde(default)]
    pub attention_bias: Option<bool>,
    #[serde(default, rename = "head_dim")]
    pub explicit_head_dim: Option<usize>,
    #[serde(default)]
    pub rope_scaling: Option<RopeScaling>,
    #[serde(default)]
    pub swiglu_limit: Option<f64>,
}

impl ModelConfig {
    pub fn from_file(path: &Path) -> Self {
        let data = std::fs::read_to_string(path)
            .unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
        serde_json::from_str(&data)
            .unwrap_or_else(|e| panic!("failed to parse {}: {e}", path.display()))
    }

    pub fn hidden(&self) -> usize {
        self.hidden_size.or(self.n_embd).expect("hidden_size or n_embd required")
    }

    pub fn num_heads(&self) -> usize {
        self.num_attention_heads.or(self.n_head).expect("num_attention_heads or n_head required")
    }

    pub fn num_layers(&self) -> usize {
        self.num_hidden_layers.or(self.n_layer).expect("num_hidden_layers or n_layer required")
    }

    pub fn max_seq_len(&self) -> usize {
        self.max_position_embeddings.or(self.n_positions).unwrap_or(2048)
    }

    pub fn ffn_hidden(&self) -> usize {
        self.intermediate_size.or(self.n_inner).unwrap_or(self.hidden() * 4)
    }

    pub fn num_kv_heads(&self) -> usize {
        self.num_key_value_heads.unwrap_or(self.num_heads())
    }

    pub fn head_dim(&self) -> usize {
        self.explicit_head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
    }

    pub fn ln_eps(&self) -> f32 {
        self.layer_norm_eps
            .or(self.layer_norm_epsilon)
            .unwrap_or(1e-5) as f32
    }

    pub fn tied_embeddings(&self) -> bool {
        self.tie_word_embeddings.unwrap_or(true)
    }

    pub fn num_experts(&self) -> usize {
        self.num_local_experts.unwrap_or(0)
    }

    pub fn experts_per_token(&self) -> usize {
        self.num_experts_per_tok.unwrap_or(1)
    }

    pub fn is_moe(&self) -> bool {
        self.num_local_experts.unwrap_or(0) > 1
    }

    pub fn is_sliding_layer(&self, layer_idx: usize) -> bool {
        self.layer_types
            .as_ref()
            .and_then(|lt| lt.get(layer_idx))
            .map(|t| t == "sliding_attention")
            .unwrap_or(false)
    }

    pub fn window_size(&self) -> usize {
        self.sliding_window.unwrap_or(0)
    }
}