brainjepa 0.2.7

Brain-JEPA fMRI Foundation Model — inference in Rust (RLX)
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

//! Minimal RLX embedding example.
//!
//! ```sh
//! cargo run --example embed_rlx --release -- \
//!   --features rlx-engine \
//!   data/brainjepa.safetensors \
//!   data/gradient_mapping_450.csv \
//!   data/test_fmri.safetensors
//! ```
use brainjepa::rlx::{ensure_device, BrainJepaEncoder};
use brainjepa::{DataConfig, ModelConfig};
use rlx::Device;

fn main() -> anyhow::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 4 {
        eprintln!(
            "usage: embed_rlx <weights> <gradient.csv> <input.safetensors> [output.safetensors]"
        );
        std::process::exit(1);
    }

    brainjepa::init_threads(None);
    let dev = Device::Cpu;
    ensure_device(dev)?;

    let (mut encoder, ms) = BrainJepaEncoder::from_weights(
        &args[1],
        &args[2],
        &ModelConfig::default(),
        &DataConfig::default(),
        &dev,
    )?;
    println!("Loaded in {ms:.0} ms: {}", encoder.describe());

    let result = encoder.encode_safetensors(&args[3])?;
    println!(
        "Encoded: {} patches x {} dims in {:.1} ms",
        result.n_patches(),
        result.embed_dim(),
        result.ms_encode
    );

    let out = args
        .get(4)
        .map(|s| s.as_str())
        .unwrap_or("embeddings.safetensors");
    result.save_safetensors(out)?;
    println!("Saved: {out}");
    Ok(())
}