alith_models/local_model/metadata/
mod.rs

1pub mod config_json;
2pub mod general;
3pub mod llm;
4pub mod tokenizer;
5use super::gguf::tools::{gguf_file::GgufFile, gguf_layers::GgufLayers};
6use general::GeneralMetadata;
7use llm::Architecture;
8use tokenizer::TokenizerMetadata;
9
10pub struct LocalLLMMetadata {
11    pub general: GeneralMetadata,
12    pub llm: Architecture,
13    pub tokenizer: TokenizerMetadata,
14    pub gguf_file: GgufFile,
15    pub layers: GgufLayers,
16}
17
18impl LocalLLMMetadata {
19    pub fn from_gguf_path(path: &std::path::Path) -> crate::Result<Self> {
20        let mut reader = std::fs::File::open(path)?;
21        let gguf: GgufFile = GgufFile::read(&mut reader)?;
22
23        Ok(Self {
24            general: GeneralMetadata::from_gguf(&gguf)?,
25            llm: Architecture::from_gguf(&gguf)?,
26            tokenizer: TokenizerMetadata::from_gguf(&gguf)?,
27            layers: GgufLayers::from_tensors(&gguf.tensors),
28            gguf_file: gguf,
29        })
30    }
31
32    pub fn estimate_model_size(&self) -> crate::Result<u64> {
33        Ok(self.gguf_file.size())
34    }
35
36    pub fn estimate_context_size(&self, ctx_size: u64, batch_size: Option<u64>) -> u64 {
37        self.llm.estimate_context_size(ctx_size, batch_size)
38    }
39
40    pub fn average_layer_size_bytes(
41        &self,
42        ctx_size: u64,
43        batch_size: Option<u64>,
44    ) -> crate::Result<u64> {
45        let total_layers_size = self.layers.total_size_blocks_bytes();
46        let block_count = self.layers.count_blocks();
47        let context_size = self.estimate_context_size(ctx_size, batch_size);
48        let total_size = total_layers_size + context_size;
49        Ok(total_size / block_count)
50    }
51
52    pub fn context_length(&self) -> u64 {
53        self.llm.context_length()
54    }
55}
56
57impl std::fmt::Debug for LocalLLMMetadata {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        let mut debug_struct = f.debug_struct("LocalLLMMetadata");
60        debug_struct.field("GeneralMetadata", &self.general);
61        debug_struct.field("Architecture", &self.llm);
62        debug_struct.field("TokenizerMetadata", &self.tokenizer);
63        debug_struct.finish()
64    }
65}