TRIBE v2 — Rust Edition
A Foundation Model of Vision, Audition, and Language for In-Silico Neuroscience
Overview
This directory contains the
same pretrained weights as
facebook/tribev2, converted to the
safetensors format for use with the pure-Rust inference engine
tribev2-rs.
No fine-tuning, quantisation, or architectural changes have been made.
The model is bit-for-bit equivalent to the original Python checkpoint — every layer has been independently verified for numerical parity.
Model description
TRIBE v2 is a deep multimodal brain encoding model that predicts fMRI responses to naturalistic stimuli (video, audio, text).
It combines three state-of-the-art feature extractors:
| Modality | Extractor | Dim |
|---|
| Text | LLaMA 3.2-3B | 3 072 |
| Audio | Wav2Vec-BERT 2.0 | 1 024 |
| Video | V-JEPA2 ViT-G | 1 408 |
These multimodal representations are projected and fused by a Transformer encoder (8 layers, 1 152-d, ScaleNorm, Rotary PE) that outputs predicted BOLD responses on the fsaverage5 cortical mesh (~20 484 vertices).
Full architectural details are in the
paper and in the
facebook/tribev2 model card.
Files
| File | Description |
|---|
model.safetensors | Pretrained weights (safetensors, converted from the original PyTorch Lightning checkpoint) |
config.yaml | Model hyper-parameters (hidden dim, depth, heads, modalities, …) |
build_args.json | Feature-extractor build arguments used at training time |
fsaverage5/ | FreeSurfer fsaverage5 cortical mesh files (.pial, .inflated, .sulc, .curv) for brain visualisation |
Encoding Input Data into Feature Tensors
The model consumes three feature tensors, one per modality, each shaped
[1, n_layers × dim, T] where T is the number of timesteps at 2 Hz
(one vector per 0.5 s).
| Modality | Extractor | Layer groups | Dim / group | Total dim |
|---|
| Text | LLaMA-3.2-3B | 2 | 3 072 | 6 144 |
| Audio | Wav2Vec-BERT 2.0 | 2 | 1 024 | 2 048 |
| Video | V-JEPA2 ViT-G | 2 | 1 408 | 2 816 |
Text — string → tensor
Text feature extraction runs entirely in Rust via
llama-cpp-rs.
Download a GGUF quantisation of
LLaMA-3.2-3B first.
Option A — raw string (uniform timing)
1use tribev2::features::{LlamaFeatureConfig, extract_llama_features, resample_features};
2use tribev2::tensor::Tensor;
3
4let config = LlamaFeatureConfig {
5 model_path: "llama-3.2-3b.gguf".into(),
6 layer_positions: vec![0.5, 0.75, 1.0], // → layers 13, 20, 27 of 28
7 n_layers: 28, // LLaMA-3.2-3B
8 n_ctx: 2048,
9 frequency: 2.0, // Hz
10};
11
12let feats = extract_llama_features(&config, "The quick brown fox", false)?;
13// feats.data: [3, 3072, n_tokens]
14
15// Resample to exactly 100 TRs and reshape to [1, 6144, 100]
16let feats = resample_features(&feats, 100);
17let text_tensor = Tensor::from_vec(
18 feats.data.data,
19 vec![1, feats.n_layers * feats.feature_dim, feats.n_timesteps],
20);
Option B — word-timed events (precise temporal alignment)
1use tribev2::features::{LlamaFeatureConfig, extract_llama_features_timed};
2
3let words = vec![
4 ("The".into(), 0.0_f64),
5 ("quick".into(), 0.3),
6 ("brown".into(), 0.55),
7 ("fox".into(), 0.82),
8];
9let total_duration = 2.0; // seconds
10
11let feats = extract_llama_features_timed(&config, &words, total_duration, false)?;
12// feats.data: [3, 3072, ceil(2.0 * 2.0) = 4]
Option C — full pipeline from a text file
1use tribev2::events::build_events_from_media;
2use tribev2::features::{LlamaFeatureConfig, extract_llama_features_timed};
3
4let events = build_events_from_media(
5 Some("transcript.txt"), // text_path
6 None, // audio_path
7 None, // video_path
8 "/tmp/cache", // cache_dir
9 "english",
10 256, // max_context_len
11)?;
12
13let words = events.words_timed(); // Vec<(String, f64)>
14let duration = events.duration();
15
16let feats = extract_llama_features_timed(&config, &words, duration, false)?;
Audio — MP3 / WAV / FLAC → tensors
Audio features come from two sources:
- Text channel — transcribe the audio → word timestamps → LLaMA
(full Rust pipeline, no Python needed)
- Audio channel — Wav2Vec-BERT 2.0 activations
(pre-extract in Python; see Pre-extracted features)
Transcribe audio → text features (Rust)
Requires whisperx or whisper (pip install whisperx) and ffmpeg.
1use tribev2::events::{transcribe_audio, build_events_from_media};
2use tribev2::features::{LlamaFeatureConfig, extract_llama_features_timed};
3
4// Option A: transcribe directly
5let events = transcribe_audio("interview.mp3", "english", 0.0)?;
6let words = events.words_timed();
7let feats = extract_llama_features_timed(&config, &words, events.duration(), false)?;
8
9// Option B: full pipeline (also attaches Audio events to the list)
10let events = build_events_from_media(
11 None,
12 Some("interview.mp3"), // audio_path
13 None,
14 "/tmp/cache", "english", 256,
15)?;
16let feats = extract_llama_features_timed(
17 &config, &events.words_timed(), events.duration(), false,
18)?;
Transcript caching — transcribe_audio saves the whisperX JSON next to
the audio file (interview.json) and reloads it on subsequent calls,
avoiding repeated transcription.
Video — MP4 → tensors
Video features come from two sources:
- Text channel — extract audio → transcribe → LLaMA (Rust)
- Video channel — V-JEPA2 ViT-G activations
(pre-extract in Python; see Pre-extracted features)
MP4 file
1use tribev2::events::build_events_from_media;
2
3let events = build_events_from_media(
4 None, None,
5 Some("clip.mp4"), // video_path
6 "/tmp/cache", "english", 256,
7)?;
8let feats = extract_llama_features_timed(
9 &config, &events.words_timed(), events.duration(), false,
10)?;
Sequence of images (PNG / JPG / WEBP / …)
Convert each frame (or the whole sequence) to an MP4 first, then use the video path above.
1use tribev2::events::create_video_from_image;
2
3// Single static image held for N seconds
4let mp4 = create_video_from_image("frame.png", 5.0, 24, "/tmp/cache")?;
5
6// Image sequence → MP4 via ffmpeg (shell out)
7std::process::Command::new("ffmpeg")
8 .args(["-y", "-framerate", "24"])
9 .args(["-pattern_type", "glob", "-i", "frames/*.png"])
10 .args(["-c:v", "libx264", "-pix_fmt", "yuv420p"])
11 .arg("/tmp/cache/sequence.mp4")
12 .status()?;
13
14let events = build_events_from_media(
15 None, None, Some("/tmp/cache/sequence.mp4"),
16 "/tmp/cache", "english", 256,
17)?;
Pre-extracted features (Python)
Wav2Vec-BERT and V-JEPA2 have no Rust implementation yet.
Extract them in Python and save as raw float32 binary files:
1import numpy as np
2from tribev2 import TribeModel
3
4model = TribeModel.from_pretrained("facebook/tribev2", cache_folder="./cache")
5df = model.get_events_dataframe(video_path="clip.mp4")
6
7# Extract features: dict {modality: np.ndarray [n_layers, dim, T]}
8features = model.extract_features(df)
9
10# Save each modality as a flat float32 binary
11for modality, arr in features.items():
12 arr.astype(np.float32).flatten().tofile(f"{modality}_features.bin")
13 print(f"{modality}: {arr.shape}") # e.g. audio: (2, 1024, 200)
Load them in Rust:
1use tribev2::tensor::Tensor;
2
3fn load_features(path: &str, n_layers: usize, dim: usize, t: usize)
4 -> anyhow::Result<Tensor>
5{
6 let bytes = std::fs::read(path)?;
7 let data: Vec<f32> = bytes.chunks_exact(4)
8 .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
9 .collect();
10 Ok(Tensor::from_vec(data, vec![1, n_layers * dim, t]))
11}
12
13// audio: 2 layer groups × 1024 dim × 200 timesteps → [1, 2048, 200]
14let audio = load_features("audio_features.bin", 2, 1024, 200)?;
15// video: 2 layer groups × 1408 dim × 200 timesteps → [1, 2816, 200]
16let video = load_features("video_features.bin", 2, 1408, 200)?;
Putting it all together
1use std::collections::BTreeMap;
2use tribev2::config::TribeV2Config;
3use tribev2::events::build_events_from_media;
4use tribev2::features::{LlamaFeatureConfig, extract_llama_features_timed, resample_features};
5use tribev2::model::tribe::TribeV2;
6use tribev2::tensor::Tensor;
7use tribev2::weights::{WeightMap, load_weights};
8
9// Load model
10let config: TribeV2Config = serde_yaml::from_str(
11 &std::fs::read_to_string("data/config.yaml")?
12)?;
13let mut model = TribeV2::new(
14 tribev2::ModelBuildArgs::from_json("data/build_args.json")?.to_modality_dims(),
15 20484, 100, &config.brain_model_config,
16);
17load_weights(
18 &mut WeightMap::from_safetensors("data/model.safetensors")?,
19 &mut model,
20)?;
21
22// 1. Build events from a video file (transcribes audio automatically)
23let events = build_events_from_media(
24 None, None, Some("clip.mp4"),
25 "/tmp/cache", "english", 256,
26)?;
27let n_trs = 100;
28
29// 2. Text features via LLaMA (Rust)
30let llama_cfg = LlamaFeatureConfig {
31 model_path: "llama-3.2-3b.gguf".into(),
32 ..Default::default()
33};
34let text_raw = extract_llama_features_timed(
35 &llama_cfg, &events.words_timed(), events.duration(), false,
36)?;
37let text_raw = resample_features(&text_raw, n_trs);
38let text = Tensor::from_vec(
39 text_raw.data.data,
40 vec![1, text_raw.n_layers * text_raw.feature_dim, n_trs],
41);
42
43// 3. Audio + video features pre-extracted in Python and saved as .bin
44let audio = load_features("audio_features.bin", 2, 1024, n_trs)?;
45let video = load_features("video_features.bin", 2, 1408, n_trs)?;
46
47// 4. Run inference → [1, 20484, 100] predicted BOLD on fsaverage5
48let mut features = BTreeMap::new();
49features.insert("text".into(), text);
50features.insert("audio".into(), audio);
51features.insert("video".into(), video);
52
53let output = model.forward(&features, None, true);
Rust usage
1use std::collections::BTreeMap;
2use tribev2::model::tribe::TribeV2;
3use tribev2::tensor::Tensor;
4
5// Load model from this data directory
6let model = TribeV2::from_pretrained(
7 "data/config.yaml",
8 "data/model.safetensors",
9 Some("data/build_args.json"),
10).unwrap();
11
12// Build multi-modal feature tensors [1, dim, T]
13let mut features = BTreeMap::new();
14features.insert("text".to_string(), Tensor::zeros(&[1, 6144, 100]));
15features.insert("audio".to_string(), Tensor::zeros(&[1, 2048, 100]));
16features.insert("video".to_string(), Tensor::zeros(&[1, 2816, 100]));
17
18// Forward pass → [1, 20484, 100]
19let output = model.forward(&features, None, true);
20println!("{:?}", output.shape()); // [1, 20484, 100]
See the
tribev2-rs README for the full CLI, feature flags, benchmarks, and brain-visualisation API.
Converting weights from the original checkpoint
1# 1. Download the original checkpoint from HuggingFace
2cargo run --bin tribev2-download --features hf-download -- --repo facebook/tribev2
3
4# 2. Convert to safetensors (requires Python ≥ 3.9, torch, safetensors)
5python3 scripts/convert_checkpoint.py weights/best.ckpt data/model.safetensors
6# → data/model.safetensors + data/build_args.json
Pretrained model parameters
| Parameter | Value |
|---|
| Hidden dim | 1 152 |
| Encoder depth | 8 |
| Attention heads | 8 |
| FF multiplier | 4× |
| Norm | ScaleNorm |
| Position encoding | Rotary (dim = 72) |
| Low-rank head | 2 048 |
| Subjects (released) | 1 (average subject) |
| Output surface | fsaverage5 (20 484 vertices) |
| Output timesteps | 100 TRs |
Citation
If you use these weights or the Rust inference engine, please cite the original paper:
1@article{dAscoli2026TribeV2,
2 title={A foundation model of vision, audition, and language for in-silico neuroscience},
3 author={d'Ascoli, St{\'e}phane and Rapin, J{\'e}r{\'e}my and Benchetrit, Yohann and
4 Brookes, Teon and Begany, Katelyn and Raugel, Jos{\'e}phine and
5 Banville, Hubert and King, Jean-R{\'e}mi},
6 year={2026}
7}
License
The
model weights (all files in this directory) are released under the
Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0) license,
identical to the original
facebook/tribev2 release.
You are free to share and adapt the weights for non-commercial purposes,
provided you give appropriate credit and indicate if changes were made.
Commercial use is not permitted.
The Rust source code of tribev2-rs is separately licensed under Apache-2.0.