rlx_models_core/
weights.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Model-agnostic weight I/O — paths, formats, drain policy only.
17//!
18//! Architecture checks and tensor renaming live in **model crates** ([`gguf_validate_arch`],
19//! [`register_gguf_tensor_resolver`]), not here.
20//!
21//! ```no_run
22//! use rlx_models_core::weights::{self, LoadOpts};
23//!
24//! let (path, map) = weights::open_map_with(LoadOpts::map().prefer_q4_k_m(), "weights/")?;
25//! let loaded = weights::open_with(LoadOpts::loader(), "model.gguf")?;
26//! # Ok::<(), anyhow::Error>(())
27//! ```
28
29use std::path::{Path, PathBuf};
30
31use anyhow::Result;
32
33pub use crate::gguf_resolve::{
34    GgufTensorNameResolver, LlamaFamilyGgufResolver, PassThroughGgufResolver,
35    PrefixStripGgufResolver, Qwen35NativeGgufResolver, register_gguf_tensor_resolver,
36};
37pub use crate::gguf_support::{
38    ResolveWeightsOptions, gguf_split_siblings, gguf_validate_arch, list_gguf_files_in_dir,
39    load_gguf_file, resolve_weights_file, resolve_weights_file_with_options,
40};
41pub use crate::weight_loader::{GgufLoader, WeightLoader};
42pub use crate::weight_map::{WeightDrainPolicy, WeightMap};
43pub use crate::weight_registry::{
44    LoadWeightsOptions, LoadedWeights, RegisteredFormat, WeightFormatRegistration,
45    format_for_extension, list_registered_formats, load_weight_map_resolved, load_weights_resolved,
46    open_weight_loader, register_weight_format,
47};
48
49/// Alias for [`LoadWeightsOptions`].
50pub type LoadOpts<'a> = LoadWeightsOptions<'a>;
51
52/// Alias for [`ResolveWeightsOptions`].
53pub type ResolveOpts<'a> = ResolveWeightsOptions<'a>;
54
55/// Default directory resolve: prefer [`DEFAULT_GGUF_PREFER_SUBSTR`](crate::gguf_support::DEFAULT_GGUF_PREFER_SUBSTR).
56pub fn default_resolve_opts<'a>() -> ResolveWeightsOptions<'a> {
57    ResolveWeightsOptions::default()
58        .prefer_substring(crate::gguf_support::DEFAULT_GGUF_PREFER_SUBSTR)
59}
60
61/// Resolve a file or weights directory to one on-disk path.
62pub fn pick(path: impl AsRef<Path>, resolve: &ResolveWeightsOptions<'_>) -> Result<PathBuf> {
63    resolve_weights_file_with_options(path.as_ref(), resolve)
64}
65
66/// [`pick`] with [`default_resolve_opts`].
67pub fn pick_default(path: impl AsRef<Path>) -> Result<PathBuf> {
68    pick(path, &default_resolve_opts())
69}
70
71/// Resolve a weights path, validate GGUF `general.architecture` when applicable, drain to F32.
72///
73/// Call from model runners with that family's arch list (e.g. [`crate::DINOV2_GGUF_ARCHES`]).
74pub fn load_weight_map(path: impl AsRef<Path>, gguf_arches: &[&str]) -> Result<WeightMap> {
75    let path = path.as_ref();
76    let file = pick_default(path)?;
77    if file.extension().and_then(|s| s.to_str()) == Some("gguf") {
78        gguf_validate_arch(&file, gguf_arches)?;
79    }
80    WeightMap::from_resolved_path(path)
81}
82
83/// Idempotent: ensure built-in GGUF tensor resolvers are registered (safe to call from `main`).
84pub fn init() {
85    crate::gguf_resolve::ensure_builtin_resolvers();
86}
87
88impl WeightFormatRegistration {
89    /// Describe a custom on-disk format.
90    pub const fn new(
91        id: &'static str,
92        extensions: &'static [&'static str],
93        open: crate::weight_registry::WeightLoaderFactory,
94    ) -> Self {
95        Self {
96            id,
97            extensions,
98            open,
99        }
100    }
101
102    /// [`register_weight_format`] shorthand.
103    pub fn register(self) {
104        register_weight_format(self);
105    }
106}
107
108/// Resolve + open (live [`WeightLoader`]).
109pub fn open(path: impl AsRef<Path>) -> Result<LoadedWeights> {
110    open_with(LoadOpts::loader(), path)
111}
112
113/// Resolve + open with options.
114pub fn open_with(opts: LoadOpts<'_>, path: impl AsRef<Path>) -> Result<LoadedWeights> {
115    load_weights_resolved(path.as_ref(), opts)
116}
117
118/// Resolve + drain to F32 [`WeightMap`].
119pub fn open_map(path: impl AsRef<Path>) -> Result<(PathBuf, WeightMap)> {
120    open_map_with(LoadOpts::map(), path)
121}
122
123/// Resolve + drain with options.
124pub fn open_map_with(opts: LoadOpts<'_>, path: impl AsRef<Path>) -> Result<(PathBuf, WeightMap)> {
125    load_weight_map_resolved(path.as_ref(), opts)
126}
127
128/// Numbered `.gguf` listing + resolve hints for a directory (CLI / errors).
129pub fn gguf_dir_guide(dir: &Path) -> Result<GgufDirGuide> {
130    let files = list_gguf_files_in_dir(dir)?;
131    let mut lines = Vec::new();
132    if files.is_empty() {
133        lines.push(format!("No .gguf files in {dir:?}"));
134        return Ok(GgufDirGuide { files, lines });
135    }
136    lines.push(format!("{} .gguf file(s) in {dir:?}:", files.len()));
137    for (i, p) in files.iter().enumerate() {
138        let name = p.file_name().and_then(|s| s.to_str()).unwrap_or("?");
139        lines.push(format!("  [{i}] {name}"));
140    }
141    if files.len() > 1 {
142        lines.push(String::new());
143        lines.push("Pick one:".into());
144        lines.push(format!("  pass exact path: {:?}", files[0]));
145        lines.push("  Rust: LoadOpts::map().prefer_q4_k_m()".into());
146        lines.push("  Rust: LoadOpts::map().gguf_index(0)".into());
147        lines.push("  CLI:  rlx-inspect dir/ --prefer Q4_K_M".into());
148    }
149    Ok(GgufDirGuide { files, lines })
150}
151
152#[derive(Debug, Clone)]
153pub struct GgufDirGuide {
154    pub files: Vec<PathBuf>,
155    pub lines: Vec<String>,
156}
157
158impl GgufDirGuide {
159    pub fn print(&self) {
160        for line in &self.lines {
161            println!("{line}");
162        }
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    #[test]
170    fn load_opts_format_only() {
171        assert!(LoadOpts::map().into_map);
172        assert!(!LoadOpts::loader().into_map);
173    }
174}
rlx_models_core/weights.rs

rlx_models_core/
weights.rs
