brainharmony 0.1.0

Brain-Harmony multimodal brain foundation model — inference in Rust with Burn ML
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

/// Minimal embedding example.
///
/// ```sh
/// cargo run --example embed --release -- \
///     data/brainharmony.safetensors \
///     data/gradient_mapping_400.csv \
///     data/schaefer400_roi_eigenmodes.csv \
///     data/test_signal.safetensors
/// ```
use brainharmony::prelude::*;
use burn::backend::NdArray;

type B = NdArray;

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

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    brainharmony::init_threads(None);

    let (encoder, ms) = BrainHarmonyEncoder::<B>::from_weights(
        &args[1], &args[2], &args[3],
        &ModelConfig::default(),
        &DataConfig::default(),
        &device,
    )?;
    println!("Loaded in {ms:.0} ms: {}", encoder.describe());

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

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

    Ok(())
}