Skip to main content

alith_models/local_model/metadata/
general.rs

1use crate::local_model::gguf::tools::gguf_tensors::GgmlDType;
2
3#[derive(Clone)]
4pub struct GeneralMetadata {
5    // Required fields
6    pub architecture: String,
7    pub quantization_version: u32,
8    pub alignment: u32,
9
10    // Optional metadata fields
11    pub name: Option<String>,
12    pub author: Option<String>,
13    pub version: Option<String>,
14    pub organization: Option<String>,
15    pub basename: Option<String>,
16    pub finetune: Option<String>,
17    pub description: Option<String>,
18    pub quantized_by: Option<String>,
19    pub size_label: Option<String>,
20    pub license: Option<String>,
21    pub license_name: Option<String>,
22    pub license_link: Option<String>,
23    pub url: Option<String>,
24    pub doi: Option<String>,
25    pub uuid: Option<String>,
26    pub repo_url: Option<String>,
27    pub tags: Option<Vec<String>>,
28    pub languages: Option<Vec<String>>,
29    pub datasets: Option<Vec<String>>,
30    pub file_type: Option<FileType>,
31
32    // Source metadata
33    pub source: SourceMetadata,
34}
35
36#[derive(Clone)]
37pub struct SourceMetadata {
38    pub url: Option<String>,
39    pub doi: Option<String>,
40    pub uuid: Option<String>,
41    pub repo_url: Option<String>,
42}
43
44#[derive(Debug, Clone, Copy)]
45pub enum FileType {
46    AllF32 = 0,
47    MostlyF16 = 1,
48    MostlyQ4_0 = 2,
49    MostlyQ4_1 = 3,
50    MostlyQ4_1SomeF16 = 4,
51    MostlyQ8_0 = 7,
52    MostlyQ5_0 = 8,
53    MostlyQ5_1 = 9,
54    MostlyQ2K = 10,
55    MostlyQ3KS = 11,
56    MostlyQ3KM = 12,
57    MostlyQ3KL = 13,
58    MostlyQ4KS = 14,
59    MostlyQ4KM = 15,
60    MostlyQ5KS = 16,
61    MostlyQ5KM = 17,
62    MostlyQ6K = 18,
63}
64
65impl FileType {
66    pub fn from_u32(u: u32) -> crate::Result<Self> {
67        match u {
68            0 => Ok(Self::AllF32),
69            1 => Ok(Self::MostlyF16),
70            2 => Ok(Self::MostlyQ4_0),
71            3 => Ok(Self::MostlyQ4_1),
72            4 => Ok(Self::MostlyQ4_1SomeF16),
73            7 => Ok(Self::MostlyQ8_0),
74            8 => Ok(Self::MostlyQ5_0),
75            9 => Ok(Self::MostlyQ5_1),
76            10 => Ok(Self::MostlyQ2K),
77            11 => Ok(Self::MostlyQ3KS),
78            12 => Ok(Self::MostlyQ3KM),
79            13 => Ok(Self::MostlyQ3KL),
80            14 => Ok(Self::MostlyQ4KS),
81            15 => Ok(Self::MostlyQ4KM),
82            16 => Ok(Self::MostlyQ5KS),
83            17 => Ok(Self::MostlyQ5KM),
84            18 => Ok(Self::MostlyQ6K),
85            _ => crate::bail!("Invalid FileType value: {}", u),
86        }
87    }
88
89    pub fn to_ggml_d_type(&self) -> GgmlDType {
90        match self {
91            Self::AllF32 => GgmlDType::F32,
92            Self::MostlyF16 => GgmlDType::F16,
93            Self::MostlyQ4_0 => GgmlDType::Q4_0,
94            Self::MostlyQ4_1 => GgmlDType::Q4_1,
95            Self::MostlyQ4_1SomeF16 => GgmlDType::Q4_1,
96            Self::MostlyQ8_0 => GgmlDType::Q8_0,
97            Self::MostlyQ5_0 => GgmlDType::Q5_0,
98            Self::MostlyQ5_1 => GgmlDType::Q5_1,
99            Self::MostlyQ2K => GgmlDType::Q2K,
100            Self::MostlyQ3KS => GgmlDType::Q3K,
101            Self::MostlyQ3KM => GgmlDType::Q3K,
102            Self::MostlyQ3KL => GgmlDType::Q3K,
103            Self::MostlyQ4KS => GgmlDType::Q4K,
104            Self::MostlyQ4KM => GgmlDType::Q4K,
105            Self::MostlyQ5KS => GgmlDType::Q5K,
106            Self::MostlyQ5KM => GgmlDType::Q5K,
107            Self::MostlyQ6K => GgmlDType::Q6K,
108        }
109    }
110}
111
112impl GeneralMetadata {
113    pub fn from_gguf(
114        gguf: &crate::local_model::gguf::tools::gguf_file::GgufFile,
115    ) -> crate::Result<Self> {
116        let file_type: Option<u32> = gguf.get_value("general.file_type")?;
117        let file_type = file_type.map(FileType::from_u32).transpose()?;
118        Ok(Self {
119            architecture: gguf.get_value("general.architecture")?,
120            quantization_version: gguf.get_value("general.quantization_version")?,
121            alignment: gguf.get_value("general.alignment")?,
122            name: gguf.get_value("general.name")?,
123            author: gguf.get_value("general.author")?,
124            version: gguf.get_value("general.version")?,
125            organization: gguf.get_value("general.organization")?,
126            basename: gguf.get_value("general.basename")?,
127            finetune: gguf.get_value("general.finetune")?,
128            description: gguf.get_value("general.description")?,
129            quantized_by: gguf.get_value("general.quantized_by")?,
130            size_label: gguf.get_value("general.size_label")?,
131            license: gguf.get_value("general.license")?,
132            license_name: gguf.get_value("general.license_name")?,
133            license_link: gguf.get_value("general.license_link")?,
134            url: gguf.get_value("general.url")?,
135            doi: gguf.get_value("general.doi")?,
136            uuid: gguf.get_value("general.uuid")?,
137            repo_url: gguf.get_value("general.repo_url")?,
138            tags: gguf.get_value("general.tags")?,
139            languages: gguf.get_value("general.languages")?,
140            datasets: gguf.get_value("general.datasets")?,
141            file_type,
142            source: SourceMetadata::from_gguf(gguf)?,
143        })
144    }
145}
146
147impl SourceMetadata {
148    pub fn from_gguf(
149        gguf: &crate::local_model::gguf::tools::gguf_file::GgufFile,
150    ) -> crate::Result<Self> {
151        Ok(Self {
152            url: gguf.get_value("general.source.url")?,
153            doi: gguf.get_value("general.source.doi")?,
154            uuid: gguf.get_value("general.source.uuid")?,
155            repo_url: gguf.get_value("general.source.repo_url")?,
156        })
157    }
158}
159
160impl std::fmt::Debug for GeneralMetadata {
161    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
162        let mut debug_struct = f.debug_struct("GeneralMetadata");
163
164        // Helper macro to add fields
165        macro_rules! add_field {
166            ($field:ident) => {
167                match &self.$field {
168                    Some(value) => debug_struct.field(stringify!($field), value),
169                    None => &mut debug_struct,
170                };
171            };
172        }
173
174        debug_struct.field("architecture", &self.architecture);
175        debug_struct.field("quantization_version", &self.quantization_version);
176        debug_struct.field("alignment", &self.alignment);
177
178        add_field!(name);
179        add_field!(author);
180        add_field!(version);
181        add_field!(organization);
182        add_field!(basename);
183        add_field!(finetune);
184        add_field!(description);
185        add_field!(quantized_by);
186        add_field!(size_label);
187        add_field!(license);
188        add_field!(license_name);
189        add_field!(license_link);
190        add_field!(url);
191        add_field!(doi);
192        add_field!(uuid);
193        add_field!(repo_url);
194        add_field!(tags);
195        add_field!(languages);
196        add_field!(datasets);
197        add_field!(file_type);
198
199        debug_struct.field("source", &self.source);
200
201        debug_struct.finish()
202    }
203}
204
205impl std::fmt::Debug for SourceMetadata {
206    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207        let mut debug_struct = f.debug_struct("SourceMetadata");
208
209        macro_rules! add_field {
210            ($field:ident) => {
211                if let Some(value) = &self.$field {
212                    debug_struct.field(stringify!($field), value);
213                }
214            };
215        }
216
217        add_field!(url);
218        add_field!(doi);
219        add_field!(uuid);
220        add_field!(repo_url);
221
222        debug_struct.finish()
223    }
224}