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

/// Batch embedding — encode multiple signal files in one call.
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: batch <weights> <gradient.csv> <geoh.csv> <input1> [input2] ...");
        std::process::exit(1);
    }

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

    let (encoder, _) = BrainHarmonyEncoder::<B>::from_weights(
        &args[1], &args[2], &args[3],
        &ModelConfig::default(),
        &DataConfig::default(),
        &device,
    )?;

    let input_paths: Vec<&str> = args[4..].iter().map(|s| s.as_str()).collect();
    println!("Encoding {} files...", input_paths.len());

    let results = encoder.encode_safetensors_batch(&input_paths)?;

    for (path, result) in input_paths.iter().zip(&results) {
        println!("  {path}: {} patches x {} dims  ({:.1} ms)",
            result.n_patches(), result.embed_dim(), result.ms_encode);
    }

    for (i, result) in results.iter().enumerate() {
        let out = format!("embeddings_{i}.safetensors");
        result.save_safetensors(&out)?;
    }
    println!("Saved {} files", results.len());

    Ok(())
}