Skip to main content

rlx_cli/
loader.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
16use crate::format::WeightFormat;
17use anyhow::{Result, anyhow};
18use rlx_core::gguf_support::{ResolveWeightsOptions, resolve_weights_file_with_options};
19use rlx_core::weight_loader::{GgufLoader, WeightLoader, hf_to_gguf_name};
20use rlx_core::weight_map::WeightMap;
21use std::path::{Path, PathBuf};
22
23pub use rlx_core::gguf_resolve::{GgufTensorNameResolver, register_gguf_tensor_resolver};
24pub use rlx_core::weight_registry::{
25    RegisteredFormat, list_registered_formats, load_weight_map_resolved, load_weights_resolved,
26    open_weight_loader,
27};
28pub use rlx_core::weights::{
29    self, GgufDirGuide, LoadOpts, ResolveOpts, gguf_dir_guide, open as open_weights, open_map,
30    open_map_with, open_with,
31};
32pub use rlx_core::{
33    LoadWeightsOptions, LoadedWeights, WeightFormatRegistration, register_weight_format,
34};
35
36pub fn open_loader(path: &Path) -> Result<Box<dyn WeightLoader>> {
37    open_loader_with_format(path, None)
38}
39
40/// Resolve a file or weights directory, then open the right loader.
41pub fn open_loader_resolved(path: &Path) -> Result<(PathBuf, Box<dyn WeightLoader>)> {
42    open_loader_resolved_with_options(path, &ResolveWeightsOptions::default())
43}
44
45pub fn open_loader_resolved_with_options(
46    path: &Path,
47    resolve: &ResolveWeightsOptions<'_>,
48) -> Result<(PathBuf, Box<dyn WeightLoader>)> {
49    let file = resolve_weights_file_with_options(path, resolve)?;
50    let loader = open_loader_with_format(&file, None)?;
51    Ok((file, loader))
52}
53
54pub fn open_loader_with_format(
55    path: &Path,
56    format: Option<WeightFormat>,
57) -> Result<Box<dyn WeightLoader>> {
58    let fmt = WeightFormat::resolve(path, format)?;
59    let path_str = path.to_str().ok_or_else(|| anyhow!("non-utf8 path"))?;
60    match fmt {
61        WeightFormat::Safetensors => Ok(Box::new(WeightMap::from_file(path_str)?)),
62        WeightFormat::Gguf => Ok(Box::new(GgufLoader::from_file(path_str)?)),
63    }
64}
65
66/// GGUF loader with optional MTP-head visibility (LM families).
67pub fn open_gguf_loader(path: &Path, include_mtp: bool) -> Result<GgufLoader> {
68    let path_str = path.to_str().ok_or_else(|| anyhow!("non-utf8 path"))?;
69    let mut loader = GgufLoader::from_file(path_str)?;
70    loader.include_mtp(include_mtp);
71    Ok(loader)
72}
73
74/// Resolved path → drained [`WeightMap`] (safetensors or GGUF).
75pub fn open_weight_map_resolved(path: &Path) -> Result<(PathBuf, WeightMap)> {
76    open_map_with(LoadOpts::map().warn_unused(), path)
77}
78
79/// Full resolved load (loader or map, per options).
80pub fn open_weights_resolved(path: &Path, opts: LoadWeightsOptions<'_>) -> Result<LoadedWeights> {
81    load_weights_resolved(path, opts)
82}
83
84pub fn debug_resolve_name(hf: &str) -> Option<String> {
85    hf_to_gguf_name(hf)
86}