tribev2 0.0.4

TRIBE v2 — multimodal fMRI brain encoding model inference in Rust
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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! Example: Burn backend inference for TRIBE v2.
//!
//! Demonstrates building and running the burn-based model with the
//! NdArray CPU backend. Shows both random-weight and pretrained paths.
//!
//! ```bash
//! # CPU (NdArray)
//! cargo run --example burn_inference
//!
//! # GPU (Metal)
//! cargo run --example burn_inference --no-default-features --features wgpu-metal,llama-metal
//! ```

use tribev2::config::*;
use tribev2::model_burn::tribe::TribeV2Burn;

// ── Backend selection ─────────────────────────────────────────────────────
#[cfg(all(feature = "ndarray", not(feature = "wgpu")))]
mod backend {
    pub type B = burn::backend::NdArray;
    pub fn device() -> <B as burn::prelude::Backend>::Device { Default::default() }
    pub const NAME: &str = "NdArray";
}

#[cfg(all(feature = "wgpu", not(feature = "ndarray")))]
mod backend {
    pub type B = burn::backend::Wgpu;
    pub fn device() -> burn::backend::wgpu::WgpuDevice {
        burn::backend::wgpu::WgpuDevice::DefaultDevice
    }
    pub const NAME: &str = "wgpu";
}

#[cfg(all(feature = "ndarray", feature = "wgpu"))]
mod backend {
    pub type B = burn::backend::NdArray;
    pub fn device() -> <B as burn::prelude::Backend>::Device { Default::default() }
    pub const NAME: &str = "NdArray (wgpu also enabled; using NdArray)";
}

use backend::{B, device, NAME};
use burn::prelude::*;

fn main() -> anyhow::Result<()> {
    println!("TRIBE v2 — Burn Backend Inference");
    println!("=================================");
    println!("Backend: {}\n", NAME);

    let dev = device();

    // ── Build model from config ───────────────────────────────────────
    let config = BrainModelConfig {
        hidden: 256,
        max_seq_len: 128,
        extractor_aggregation: "cat".into(),
        layer_aggregation: "cat".into(),
        linear_baseline: false,
        time_pos_embedding: true,
        subject_embedding: false,
        low_rank_head: Some(64),
        combiner: None,
        temporal_smoothing: None,
        projector: Default::default(),
        encoder: Some(EncoderConfig {
            heads: 4,
            depth: 2,
            ff_mult: 4,
            ..Default::default()
        }),
        subject_layers: Some(SubjectLayersConfig {
            n_subjects: 0,
            bias: true,
            subject_dropout: Some(0.1),
            average_subjects: true,
            ..Default::default()
        }),
        ..Default::default()
    };

    let feature_dims = vec![
        ModalityDims::new("text", 2, 128),
        ModalityDims::new("audio", 2, 64),
    ];
    let n_outputs = 500;
    let n_output_timesteps = 20;

    let model = TribeV2Burn::<B>::new(&feature_dims, n_outputs, n_output_timesteps, &config, &dev);

    println!("Model built:");
    println!("  Hidden: {}", config.hidden);
    println!("  Encoder depth: {}", config.encoder.as_ref().map_or(0, |e| e.depth));
    println!("  Low-rank head: {:?}", config.low_rank_head);
    println!("  Output: {} vertices × {} timesteps", n_outputs, n_output_timesteps);

    // ── Create synthetic features ─────────────────────────────────────
    let t = 40;
    let text  = Tensor::<B, 3>::zeros([1, 2 * 128, t], &dev);
    let audio = Tensor::<B, 3>::zeros([1, 2 * 64, t], &dev);

    let features = vec![("text", text), ("audio", audio)];

    // ── Forward pass ──────────────────────────────────────────────────
    let t0 = std::time::Instant::now();
    let output = model.forward(features);
    let elapsed = t0.elapsed();

    let [b, d, t_out] = output.dims();
    println!("\nForward pass:");
    println!("  Output: [{}, {}, {}]", b, d, t_out);
    println!("  Time: {:.1} ms", elapsed.as_secs_f64() * 1000.0);

    // ── Load pretrained weights (if available) ────────────────────────
    // Uncomment to load real weights:
    //
    // use tribev2::model_burn::weights::{BurnWeightStore, load_burn_weights};
    // let mut ws = BurnWeightStore::from_safetensors("data/model.safetensors")?;
    // let mut model = TribeV2Burn::<B>::new(&ModalityDims::pretrained(), 20484, 100,
    //     &pretrained_config, &dev);
    // load_burn_weights(&mut ws, &mut model, &dev)?;

    println!("\nDone!");
    Ok(())
}