TRIBE v2

To train a model from scratch, follow these steps:

1. Set environment variables

Configure your data and output paths using environment variables (or edit tribev2/grids/defaults.py directly):

• DATAPATH: Path to your studies.
• SAVEPATH: Path to your output.

2. Run training

• Local test run: Use the test run module.
• Grid search on Slurm: Use the cortical or subcortical run modules.

# 1. Set environment variables
export DATAPATH="/path/to/studies"
export SAVEPATH="/path/to/output"

# 2. Run training
# Local test run
python -m tribev2.grids.test_run

# Grid search on Slurm
python -m tribev2.grids.run_cortical
python -m tribev2.grids.run_subcortical

You can install TRIBE v2 in different modes depending on your requirements:

• Basic (Inference only): Installs the core package.
• With brain visualization: Adds dependencies for plotting.
• With training dependencies: Adds dependencies for training (e.g., PyTorch Lightning, W&B).

The TribeModel class provides a high-level wrapper for generating fMRI-like brain-activity predictions from text, audio, or video inputs. It follows a from_pretrained / predict pattern.

Typical workflow:

1. Load a model using from_pretrained (from a local directory or HuggingFace Hub).
2. Generate an events DataFrame using get_events_dataframe by providing a path to a text, audio, or video file.
3. Run inference using predict to get brain activity predictions and corresponding segments.

model = TribeModel.from_pretrained("facebook/tribev2")
events = model.get_events_dataframe(video_path="clip.mp4")
preds, segments = model.predict(events)

To install TRIBE v2 with plotting support, use uv pip to install directly from the GitHub repository. If using Google Colab, ensure you have activated a GPU runtime before installation and restart your environment after the installation completes.

!uv pip install "tribev2[plotting] @ git+https://github.com/facebookresearch/tribev2.git"

Use TribeModel.from_pretrained to load a model from HuggingFace. You can then generate an events dataframe from video, audio, or text and pass it to model.predict to get fMRI brain response predictions.

Key Details:

• Predictions are for the "average" subject on the fsaverage5 cortical mesh (~20k vertices).
• Predictions are offset by 5 seconds in the past to compensate for hemodynamic lag.
• get_events_dataframe accepts video_path, text_path, or audio_path. Text is automatically converted to speech and transcribed for word-level timings.

from tribev2 import TribeModel

model = TribeModel.from_pretrained("facebook/tribev2", cache_folder="./cache")

df = model.get_events_dataframe(video_path="path/to/video.mp4")
preds, segments = model.predict(events=df)
print(preds.shape)  # (n_timesteps, n_vertices)

The TemporalSmoothing class provides a 1D convolutional layer with a Gaussian kernel for temporal smoothing.

Configuration options:

• kernel_size: The size of the smoothing kernel (default: 9).
• sigma: The standard deviation of the Gaussian kernel. If None, the kernel is not initialized with Gaussian weights and remains a standard convolution.

To process a video, use model.get_events_dataframe(video_path=...). This method automatically extracts audio, transcribes speech using WhisperX, and prepares the multimodal features (visual, audio, and text) required for prediction. The resulting dataframe contains columns such as type, start, duration, filepath, text, and context.

video_path = CACHE_FOLDER / "sample_video.mp4"
df = model.get_events_dataframe(video_path=video_path)

The FmriEncoderModel is a torch.nn.Module that processes SegmentData batches.

forward(batch: SegmentData, pool_outputs: bool = True) -> torch.Tensor

• batch: A SegmentData object containing the input tensors (e.g., modality data and subject_id).
• pool_outputs: If True, applies AdaptiveAvgPool1d to the output to reduce the temporal dimension to n_output_timesteps. If False, returns the full temporal sequence.

Returns a tensor of shape [Batch, Outputs, Timesteps] (if pooled) or [Batch, Outputs, Original_Timesteps] (if not pooled).

The SplitEvents class is an EventsTransform used to partition event data into training and validation sets. It uses a val_ratio to determine the split and relies on a predefined mapping of studies to specific attributes (like chunk, task, story, etc.) to ensure splits are consistent within a study.

Note: The splitting logic is study-dependent. The assign_splits function looks up the study name in a hardcoded SPLIT_ATTRIBUTES dictionary to find the column used for splitting.

The CreateVideosFromImages class transforms Image events into Video events by generating .mp4 files from the image paths.

Configuration:

• fps: The frames per second for the generated video (default: 10).
• remove_images: If True, the original Image events are removed from the DataFrame after video creation (default: True).
• infra: Uses exca.MapInfra for parallel processing (default: exca.MapInfra(cluster="processpool")).

Dependencies:

• Requires moviepy for video generation.