model_ref/
lib.rs

1use std::{error::Error, fmt, path::Path, str::FromStr};
2
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
6pub struct ModelRef {
7    pub repo: String,
8    pub revision: Option<String>,
9    pub selector: Option<String>,
10}
11
12impl ModelRef {
13    pub fn parse(input: &str) -> Result<Self, ModelRefParseError> {
14        parse_model_ref(input)
15    }
16
17    pub fn display_id(&self) -> String {
18        format_model_ref(
19            &self.repo,
20            self.revision.as_deref(),
21            self.selector.as_deref(),
22        )
23    }
24}
25
26impl fmt::Display for ModelRef {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        f.write_str(&self.display_id())
29    }
30}
31
32impl FromStr for ModelRef {
33    type Err = ModelRefParseError;
34
35    fn from_str(input: &str) -> Result<Self, Self::Err> {
36        Self::parse(input)
37    }
38}
39
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub struct ModelRefParseError {
42    input: String,
43}
44
45impl ModelRefParseError {
46    fn new(input: &str) -> Self {
47        Self {
48            input: input.to_string(),
49        }
50    }
51}
52
53impl fmt::Display for ModelRefParseError {
54    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
55        write!(
56            f,
57            "expected model ref like org/repo, org/repo:Q4_K_M, or org/repo@rev:Q4_K_M: {}",
58            self.input
59        )
60    }
61}
62
63impl Error for ModelRefParseError {}
64
65pub fn parse_model_ref(input: &str) -> Result<ModelRef, ModelRefParseError> {
66    let trimmed = input.trim();
67    if trimmed.is_empty() {
68        return Err(ModelRefParseError::new(input));
69    }
70
71    if let Some(parsed) = parse_huggingface_repo_url(trimmed) {
72        return Ok(parsed);
73    }
74
75    parse_huggingface_repo_ref(trimmed).ok_or_else(|| ModelRefParseError::new(input))
76}
77
78pub fn format_model_ref(repo: &str, revision: Option<&str>, selector: Option<&str>) -> String {
79    match (revision, selector) {
80        (Some(revision), Some(selector)) => format!("{repo}@{revision}:{selector}"),
81        (Some(revision), None) => format!("{repo}@{revision}"),
82        (None, Some(selector)) => format!("{repo}:{selector}"),
83        (None, None) => repo.to_string(),
84    }
85}
86
87pub fn format_gguf_selection_ref(repo: &str, file: &str, selector: &str) -> String {
88    let directory_selector = Path::new(file)
89        .parent()
90        .and_then(Path::to_str)
91        .filter(|directory| !directory.is_empty() && *directory != ".");
92    format_model_ref(repo, None, directory_selector.or(Some(selector)))
93}
94
95pub fn format_canonical_ref(repo: &str, revision: &str, file: &str) -> String {
96    format!("{repo}@{revision}/{file}")
97}
98
99pub fn quant_selector_from_gguf_file(file: &str) -> Option<String> {
100    if !file.ends_with(".gguf") {
101        return None;
102    }
103
104    if let Some((prefix, _)) = file.split_once('/')
105        && is_quant_like_selector(prefix)
106    {
107        return Some(prefix.to_string());
108    }
109
110    let basename = Path::new(file).file_name()?.to_str()?;
111    let mut stem = basename.strip_suffix(".gguf")?;
112    if let Some(prefix) = split_gguf_shard_stem_prefix(stem) {
113        stem = prefix;
114    }
115
116    // Markers are matched case-insensitively. Real-world GGUF filenames
117    // are inconsistent: `Qwen3-32B-Q4_K_M.gguf` (uppercase markers) is
118    // common, but so is `qwen2.5-3b-instruct-q4_k_m.gguf` (all
119    // lowercase). We lowercase to find the marker position, but the
120    // returned selector is sliced out of the *original* stem so the
121    // model's preferred display casing survives into the public ID.
122    // The downstream matcher in `gguf_matches_quant_selector` lowercases
123    // both sides at comparison time, so preserving casing here is safe
124    // for matching while keeping IDs round-trippable.
125    let stem_lower = stem.to_ascii_lowercase();
126    for marker in [
127        "-ud-", ".ud-", "-iq", ".iq", "-q", ".q", "-bf16", ".bf16", "-f16", ".f16", "-f32", ".f32",
128    ] {
129        if let Some(pos) = stem_lower.rfind(marker) {
130            // Slice from the original (case-preserving) stem.
131            return Some(stem[pos + 1..].to_string());
132        }
133    }
134    None
135}
136
137pub fn normalize_gguf_distribution_id(file: &str) -> Option<String> {
138    let basename = Path::new(file).file_name()?.to_str()?;
139    let stem = basename.strip_suffix(".gguf")?;
140    let stem = split_gguf_shard_stem_prefix(stem).unwrap_or(stem);
141    (!stem.is_empty()).then(|| stem.to_string())
142}
143
144pub fn is_quant_like_selector(value: &str) -> bool {
145    let upper = value.to_ascii_uppercase();
146    upper.starts_with("UD-")
147        || upper.starts_with('Q')
148        || upper.starts_with("IQ")
149        || upper == "BF16"
150        || upper == "F16"
151        || upper == "F32"
152}
153
154pub fn gguf_matches_quant_selector(file: &str, selector: &str) -> bool {
155    let file_lower = file.to_ascii_lowercase();
156    let selector_lower = selector.to_ascii_lowercase();
157    if !file_lower.ends_with(".gguf") || selector_lower.is_empty() {
158        return false;
159    }
160    file_lower.contains(&format!("/{selector_lower}/"))
161        || file_lower.contains(&format!("-{selector_lower}-"))
162        || file_lower.contains(&format!(".{selector_lower}-"))
163        || file_lower.ends_with(&format!("-{selector_lower}.gguf"))
164        || file_lower.ends_with(&format!(".{selector_lower}.gguf"))
165        || file_lower.ends_with(&format!("/{selector_lower}.gguf"))
166}
167
168pub fn split_gguf_shard_info(file: &str) -> Option<SplitGgufShard<'_>> {
169    let stem = file.strip_suffix(".gguf")?;
170    let (prefix_and_part, total) = stem.rsplit_once("-of-")?;
171    if total.len() != 5 || !total.bytes().all(|byte| byte.is_ascii_digit()) {
172        return None;
173    }
174    let (prefix, part) = prefix_and_part.rsplit_once('-')?;
175    if part.len() != 5 || !part.bytes().all(|byte| byte.is_ascii_digit()) {
176        return None;
177    }
178    Some(SplitGgufShard {
179        prefix,
180        part,
181        total,
182    })
183}
184
185#[derive(Debug, Clone, Copy, PartialEq, Eq)]
186pub struct SplitGgufShard<'a> {
187    pub prefix: &'a str,
188    pub part: &'a str,
189    pub total: &'a str,
190}
191
192fn split_gguf_shard_stem_prefix(stem: &str) -> Option<&str> {
193    let (prefix_and_part, total) = stem.rsplit_once("-of-")?;
194    if total.len() != 5 || !total.bytes().all(|byte| byte.is_ascii_digit()) {
195        return None;
196    }
197    let (prefix, part) = prefix_and_part.rsplit_once('-')?;
198    if part.len() != 5 || !part.bytes().all(|byte| byte.is_ascii_digit()) {
199        return None;
200    }
201    Some(prefix)
202}
203
204fn parse_huggingface_repo_ref(input: &str) -> Option<ModelRef> {
205    let parts: Vec<&str> = input.splitn(2, '/').collect();
206    if parts.len() != 2 {
207        return None;
208    }
209    if parts[0].is_empty() || parts[1].is_empty() || parts[0].contains(':') {
210        return None;
211    }
212    let (repo_tail, revision, selector) = parse_repo_tail_selector_and_revision(parts[1])?;
213    if repo_tail.contains('/') {
214        return None;
215    }
216    Some(ModelRef {
217        repo: format!("{}/{}", parts[0], repo_tail),
218        revision,
219        selector,
220    })
221}
222
223fn parse_huggingface_repo_url(input: &str) -> Option<ModelRef> {
224    let tail = input
225        .strip_prefix("https://huggingface.co/")
226        .or_else(|| input.strip_prefix("http://huggingface.co/"))?;
227    let clean = tail
228        .split_once('?')
229        .map(|(left, _)| left)
230        .unwrap_or(tail)
231        .split_once('#')
232        .map(|(left, _)| left)
233        .unwrap_or(tail)
234        .trim_matches('/');
235    let parts: Vec<&str> = clean.split('/').collect();
236    if parts.len() < 2 || parts[0].is_empty() || parts[1].is_empty() {
237        return None;
238    }
239    let (repo_tail, revision, selector) = parse_repo_tail_selector_and_revision(parts[1])?;
240    let repo = format!("{}/{}", parts[0], repo_tail);
241    if parts.len() >= 4 && parts[2] == "tree" && !parts[3].is_empty() {
242        return Some(ModelRef {
243            repo,
244            revision: Some(parts[3].to_string()),
245            selector,
246        });
247    }
248    if parts.len() == 2 {
249        return Some(ModelRef {
250            repo,
251            revision,
252            selector,
253        });
254    }
255    None
256}
257
258fn parse_repo_tail_selector_and_revision(
259    tail: &str,
260) -> Option<(String, Option<String>, Option<String>)> {
261    let at_pos = tail.find('@');
262    let colon_pos = tail.find(':');
263
264    match (at_pos, colon_pos) {
265        (Some(at), Some(colon)) if at < colon => {
266            let repo_tail = &tail[..at];
267            let revision = &tail[at + 1..colon];
268            let selector = &tail[colon + 1..];
269            nonempty_tail(repo_tail, Some(revision), Some(selector))
270        }
271        (Some(at), Some(colon)) if colon < at => {
272            let repo_tail = &tail[..colon];
273            let selector = &tail[colon + 1..at];
274            let revision = &tail[at + 1..];
275            nonempty_tail(repo_tail, Some(revision), Some(selector))
276        }
277        (Some(at), None) => {
278            let repo_tail = &tail[..at];
279            let revision = &tail[at + 1..];
280            nonempty_tail(repo_tail, Some(revision), None)
281        }
282        (None, Some(colon)) => {
283            let repo_tail = &tail[..colon];
284            let selector = &tail[colon + 1..];
285            nonempty_tail(repo_tail, None, Some(selector))
286        }
287        (None, None) if !tail.is_empty() => Some((tail.to_string(), None, None)),
288        _ => None,
289    }
290}
291
292fn nonempty_tail(
293    repo_tail: &str,
294    revision: Option<&str>,
295    selector: Option<&str>,
296) -> Option<(String, Option<String>, Option<String>)> {
297    if repo_tail.is_empty()
298        || revision.is_some_and(str::is_empty)
299        || selector.is_some_and(str::is_empty)
300    {
301        return None;
302    }
303    Some((
304        repo_tail.to_string(),
305        revision.map(str::to_string),
306        selector.map(str::to_string),
307    ))
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313
314    #[test]
315    fn parses_public_model_refs() {
316        assert_eq!(
317            parse_model_ref("org/repo:Q4_K_M").unwrap(),
318            ModelRef {
319                repo: "org/repo".to_string(),
320                revision: None,
321                selector: Some("Q4_K_M".to_string()),
322            }
323        );
324
325        assert_eq!(
326            parse_model_ref("org/repo@abc123:Q4_K_M")
327                .unwrap()
328                .display_id(),
329            "org/repo@abc123:Q4_K_M"
330        );
331
332        assert_eq!(
333            parse_model_ref("org/repo:Q4_K_M@abc123")
334                .unwrap()
335                .display_id(),
336            "org/repo@abc123:Q4_K_M"
337        );
338    }
339
340    #[test]
341    fn rejects_mesh_exact_file_refs() {
342        assert!(parse_model_ref("org/repo/file.gguf").is_err());
343        assert!(parse_model_ref("org/repo/model.safetensors").is_err());
344        assert!(parse_model_ref("https://huggingface.co/org/repo/resolve/main/file.gguf").is_err());
345    }
346
347    #[test]
348    fn parses_huggingface_repo_urls() {
349        assert_eq!(
350            parse_model_ref("https://huggingface.co/org/repo:BF16")
351                .unwrap()
352                .display_id(),
353            "org/repo:BF16"
354        );
355        assert_eq!(
356            parse_model_ref("https://huggingface.co/org/repo/tree/main")
357                .unwrap()
358                .display_id(),
359            "org/repo@main"
360        );
361    }
362
363    #[test]
364    fn rejects_empty_or_non_repo_refs() {
365        assert!(parse_model_ref("").is_err());
366        assert!(parse_model_ref("repo-only").is_err());
367        assert!(parse_model_ref("org/:Q4_K_M").is_err());
368        assert!(parse_model_ref("org/repo:").is_err());
369    }
370
371    #[test]
372    fn extracts_quant_selectors_from_gguf_files() {
373        assert_eq!(
374            quant_selector_from_gguf_file("gemma-4-31B-it-UD-Q4_K_XL.gguf"),
375            Some("UD-Q4_K_XL".to_string())
376        );
377        assert_eq!(
378            quant_selector_from_gguf_file("Meta-Llama-3.1-8B-Instruct.Q4_K_M.gguf"),
379            Some("Q4_K_M".to_string())
380        );
381        assert_eq!(
382            quant_selector_from_gguf_file("BF16/gemma-4-31B-it-BF16-00001-of-00002.gguf"),
383            Some("BF16".to_string())
384        );
385        assert_eq!(
386            quant_selector_from_gguf_file("qwen3.5-moe-0.87B-d0.8B.Q2_K.gguf"),
387            Some("Q2_K".to_string())
388        );
389        assert_eq!(
390            quant_selector_from_gguf_file("gemma-4-31B-it-Q4_0.gguf"),
391            Some("Q4_0".to_string())
392        );
393    }
394
395    #[test]
396    fn normalizes_distribution_ids() {
397        assert_eq!(
398            normalize_gguf_distribution_id("GLM-5.1-UD-IQ2_M-00001-of-00006.gguf"),
399            Some("GLM-5.1-UD-IQ2_M".to_string())
400        );
401        assert_eq!(
402            normalize_gguf_distribution_id("Qwen3-30B-A3B-Q4_K_M.gguf"),
403            Some("Qwen3-30B-A3B-Q4_K_M".to_string())
404        );
405        assert_eq!(
406            normalize_gguf_distribution_id("UD-IQ2_M/GLM-5.1-UD-IQ2_M-00001-of-00006.gguf"),
407            Some("GLM-5.1-UD-IQ2_M".to_string())
408        );
409        assert_eq!(normalize_gguf_distribution_id("README.md"), None);
410    }
411
412    #[test]
413    fn matches_quant_selectors_against_gguf_paths() {
414        assert!(gguf_matches_quant_selector(
415            "UD-Q4_K_XL/gemma-4-31B-it-UD-Q4_K_XL-00001-of-00004.gguf",
416            "UD-Q4_K_XL"
417        ));
418        assert!(gguf_matches_quant_selector(
419            "Meta-Llama-3.1-8B-Instruct.Q4_K_M.gguf",
420            "Q4_K_M"
421        ));
422        assert!(!gguf_matches_quant_selector(
423            "Meta-Llama-3.1-8B-Instruct.Q5_K_M.gguf",
424            "Q4_K_M"
425        ));
426    }
427
428    #[test]
429    fn formats_canonical_refs() {
430        assert_eq!(
431            format_canonical_ref("org/repo", "abc123", "model.gguf"),
432            "org/repo@abc123/model.gguf"
433        );
434    }
435
436    #[test]
437    fn formats_selected_gguf_refs() {
438        assert_eq!(
439            format_gguf_selection_ref("unsloth/Qwen3-8B-GGUF", "Qwen3-8B-Q4_K_M.gguf", "Q4_K_M"),
440            "unsloth/Qwen3-8B-GGUF:Q4_K_M"
441        );
442        assert_eq!(
443            format_gguf_selection_ref(
444                "unsloth/LTX-2.3-GGUF",
445                "distilled-1.1/LTX-2.3-UD-Q4_K_M.gguf",
446                "UD-Q4_K_M"
447            ),
448            "unsloth/LTX-2.3-GGUF:distilled-1.1"
449        );
450    }
451
452    #[test]
453    fn parses_split_gguf_shard_info() {
454        assert_eq!(
455            split_gguf_shard_info("Qwen3-30B-A3B-Q4_K_M-00001-of-00004.gguf"),
456            Some(SplitGgufShard {
457                prefix: "Qwen3-30B-A3B-Q4_K_M",
458                part: "00001",
459                total: "00004",
460            })
461        );
462        assert_eq!(split_gguf_shard_info("Qwen3-30B-A3B-Q4_K_M.gguf"), None);
463    }
464}
model_ref/lib.rs

model_ref/
lib.rs
