lumamba 0.0.1

LuMamba EEG foundation model — inference in Rust on the RLX runtime
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

// lumamba-rs — LuMamba EEG foundation-model inference on the RLX runtime.
// Copyright (C) 2026 Nataliya Kosmyna.
// SPDX-License-Identifier: GPL-3.0-only

//! HuggingFace Hub weight resolution.
//!
//! [`resolve`] turns a weights/config spec into a local path: an
//! `hf://<owner>/<repo>/<file>` spec is fetched from the Hub (cached under
//! `~/.cache/huggingface`), anything else is treated as a local path. The
//! actual download lives behind the **`hf-download`** Cargo feature; without
//! it, an `hf://` spec returns a clear error pointing at the flag while local
//! paths still pass through.

use std::path::PathBuf;

use anyhow::Result;

/// Default HuggingFace repo for LuMamba weights.
pub const DEFAULT_REPO: &str = "PulpBio/LuMamba";

/// The released pretrained checkpoints in [`DEFAULT_REPO`].
pub const CHECKPOINTS: &[&str] = &[
    "LuMamba_LeJEPA_reconstruction_128slices.safetensors",
    "LuMamba_LeJEPA_reconstruction_300slices.safetensors",
    "LuMamba_LeJEPAOnly_128slices.safetensors",
    "LuMamba_ReconstructionOnly.safetensors",
];

/// Resolve a weights/config spec to a local path.
///
/// * `hf://<owner>/<repo>/<file>` → download (cached) via the Hub.
/// * anything else → returned as-is (local path).
pub fn resolve(spec: &str) -> Result<PathBuf> {
    if let Some(rest) = spec.strip_prefix("hf://") {
        let mut it = rest.splitn(3, '/');
        match (it.next(), it.next(), it.next()) {
            (Some(owner), Some(repo), Some(file))
                if !owner.is_empty() && !repo.is_empty() && !file.is_empty() =>
            {
                fetch(&format!("{owner}/{repo}"), file)
            }
            _ => anyhow::bail!(
                "bad HF spec '{spec}', expected hf://<owner>/<repo>/<file.safetensors>"
            ),
        }
    } else {
        Ok(PathBuf::from(spec))
    }
}

/// Download `file` from `repo` on the Hub, returning the cached local path.
/// Requires the `hf-download` feature.
#[cfg(feature = "hf-download")]
pub fn download(repo: &str, file: &str) -> Result<PathBuf> {
    use hf_hub::api::sync::Api;
    let api = Api::new()?;
    Ok(api.model(repo.to_string()).get(file)?)
}

#[cfg(feature = "hf-download")]
fn fetch(repo: &str, file: &str) -> Result<PathBuf> {
    eprintln!("Fetching {repo} / {file} from HuggingFace …");
    download(repo, file)
}

#[cfg(not(feature = "hf-download"))]
fn fetch(repo: &str, file: &str) -> Result<PathBuf> {
    anyhow::bail!(
        "hf:// weights ({repo}/{file}) require building with `--features hf-download`"
    )
}