ggml_rs/
lib.rs

1use regex::Regex;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5#[derive(Serialize, Deserialize, Debug, PartialEq)]
6pub struct GGUFComponents {
7    pub base_name: Option<String>,
8    pub size_label: Option<String>,
9    pub fine_tune: Option<String>,
10    pub version: String,
11    pub encoding: Option<String>,
12    pub type_field: Option<String>,
13    pub shard: Option<String>,
14}
15
16pub fn parse_gguf_filename(filename: &str) -> Option<GGUFComponents> {
17    let gguf_regex = Regex::new(
18        r"^(?P<BaseName>[A-Za-z0-9\s]*(?:(?:-(?:(?:[A-Za-z\s][A-Za-z0-9\s]*)|(?:[0-9\s]*)))*))-(?:(?P<SizeLabel>(?:\d+x)?(?:\d+\.)?\d+[A-Za-z](?:-[A-Za-z]+(\d+\.)?\d+[A-Za-z]+)?)(?:-(?P<FineTune>[A-Za-z0-9\s-]+))?)?-(?:(?P<Version>v\d+(?:\.\d+)*))(?:-(?P<Encoding>[\w_]+))?(?:-(?P<Type>LoRA|vocab))?(?:-(?P<Shard>\d{5}-of-\d{5}))?\.gguf$"
19    ).expect("Failed to compile regex");
20
21    match gguf_regex.captures(filename) {
22        Some(caps) => {
23            let mut encoding = caps.name("Encoding").map(|m| m.as_str().to_string());
24            let type_field = caps.name("Type").map(|m| m.as_str().to_string());
25
26            // Post-process to handle LoRA and vocab exclusion from Encoding
27            if let Some(enc) = encoding.clone() {
28                if enc == "LoRA" || enc == "vocab" {
29                    encoding = None;
30                }
31            }
32
33            Some(GGUFComponents {
34                base_name: caps.name("BaseName").map(|m| m.as_str().to_string()),
35                size_label: caps.name("SizeLabel").map(|m| m.as_str().to_string()),
36                fine_tune: caps.name("FineTune").map(|m| m.as_str().to_string()),
37                version: caps
38                    .name("Version")
39                    .map_or("v1.0".to_string(), |m| m.as_str().to_string()),
40                encoding,
41                type_field,
42                shard: caps.name("Shard").map(|m| m.as_str().to_string()),
43            })
44        }
45        None => None,
46    }
47}
48
49pub fn parse_gguf_filename_map(filename: &str) -> Option<HashMap<String, Option<String>>> {
50    parse_gguf_filename(filename).map(|components| {
51        let mut map = HashMap::new();
52        map.insert("BaseName".to_string(), components.base_name);
53        map.insert("SizeLabel".to_string(), components.size_label);
54        map.insert("FineTune".to_string(), components.fine_tune);
55        map.insert("Version".to_string(), Some(components.version));
56        map.insert("Encoding".to_string(), components.encoding);
57        map.insert("Type".to_string(), components.type_field);
58        map.insert("Shard".to_string(), components.shard);
59        map
60    })
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use serde_json::json;
67
68    #[test]
69    fn test_parse_gguf_filename() {
70        let test_cases = vec![
71            (
72                "Mixtral-8x7B-v0.1-KQ2.gguf",
73                json!({
74                    "BaseName": "Mixtral",
75                    "SizeLabel": "8x7B",
76                    "FineTune": null,
77                    "Version": "v0.1",
78                    "Encoding": "KQ2",
79                    "Type": null,
80                    "Shard": null
81                }),
82            ),
83            (
84                "Grok-100B-v1.0-Q4_0-00003-of-00009.gguf",
85                json!({
86                    "BaseName": "Grok",
87                    "SizeLabel": "100B",
88                    "FineTune": null,
89                    "Version": "v1.0",
90                    "Encoding": "Q4_0",
91                    "Type": null,
92                    "Shard": "00003-of-00009"
93                }),
94            ),
95            (
96                "Hermes-2-Pro-Llama-3-8B-v1.0-F16.gguf",
97                json!({
98                    "BaseName": "Hermes-2-Pro-Llama-3",
99                    "SizeLabel": "8B",
100                    "FineTune": null,
101                    "Version": "v1.0",
102                    "Encoding": "F16",
103                    "Type": null,
104                    "Shard": null
105                }),
106            ),
107            (
108                "Phi-3-mini-3.8B-ContextLength4k-instruct-v1.0.gguf",
109                json!({
110                    "BaseName": "Phi-3-mini",
111                    "SizeLabel": "3.8B-ContextLength4k",
112                    "FineTune": "instruct",
113                    "Version": "v1.0",
114                    "Encoding": null,
115                    "Type": null,
116                    "Shard": null
117                }),
118            ),
119            ("not-a-known-arrangement.gguf", json!(null)),
120        ];
121
122        for (filename, expected) in test_cases {
123            let result = parse_gguf_filename_map(filename);
124            let result_json = json!(result);
125            assert_eq!(
126                result_json, expected,
127                "Failed for filename: {}. Expected: {:?}, Got: {:?}",
128                filename, expected, result_json
129            );
130        }
131    }
132}
ggml_rs/lib.rs

ggml_rs/
lib.rs
