Skip to main content

alith_models/local_model/metadata/
config_json.rs

1use crate::local_model::gguf::memory::estimate_context_size;
2use serde::Deserialize;
3use std::{fs::File, io::BufReader};
4
5#[derive(Debug, Deserialize, Clone)]
6pub struct ConfigJson {
7    #[serde(alias = "max_position_embeddings")]
8    #[serde(alias = "n_ctx")]
9    pub context_length: u64,
10    #[serde(alias = "hidden_size")]
11    #[serde(alias = "n_embd")]
12    pub embedding_length: u64,
13    #[serde(alias = "intermediate_size")]
14    #[serde(alias = "n_inner")]
15    pub feed_forward_length: Option<u64>,
16    #[serde(alias = "num_attention_heads")]
17    #[serde(alias = "n_head")]
18    pub head_count: u64,
19    #[serde(alias = "num_key_value_heads")]
20    pub head_count_kv: Option<u64>,
21    #[serde(alias = "num_hidden_layers")]
22    #[serde(alias = "n_layers")]
23    #[serde(alias = "n_layer")]
24    #[serde(alias = "num_layers")]
25    pub block_count: u64,
26    pub torch_dtype: String,
27    pub vocab_size: u32,
28    #[serde(alias = "model_type")]
29    pub architecture: String,
30    pub model_size_bytes: Option<u64>,
31}
32
33impl ConfigJson {
34    pub fn from_local_path(config_json_path: &std::path::PathBuf) -> crate::Result<Self> {
35        let file = File::open(config_json_path)?;
36        let reader = BufReader::new(file);
37        let config: ConfigJson = serde_json::from_reader(reader)?;
38        Ok(config)
39    }
40
41    // This is converted from https://github.com/pandora-s-git/LLMVRAMCalculator/blob/main/LLMVRAMCalculator/LLMVRAMCalculator.py
42    pub fn estimate_context_size(&self, ctx_size: u64) -> u64 {
43        estimate_context_size(
44            ctx_size,
45            self.embedding_length,
46            self.head_count,
47            self.head_count_kv.unwrap(),
48            self.block_count,
49            None,
50        )
51    }
52}