yautungsky/braindance

Python

0

5 commits

updated Apr 12, 2026

See the code

README

TRIBE v2

A Foundation Model of Vision, Audition, and Language for In-Silico Neuroscience

Open In Colab License: CC BY-NC 4.0 Python 3.10+

📄 Paper ▶️ Demo | 🤗 Weights

TRIBE v2 is a deep multimodal brain encoding model that predicts fMRI brain responses to naturalistic stimuli (video, audio, text). It combines state-of-the-art feature extractors — LLaMA 3.2 (text), V-JEPA2 (video), and Wav2Vec-BERT (audio) — into a unified Transformer architecture that maps multimodal representations onto the cortical surface.

Quick start

Load a pretrained model from HuggingFace and predict brain responses to a video:

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)

Predictions are for the "average" subject (see paper for details) and live on the fsaverage5 cortical mesh (~20k vertices). You can also pass text_path or audio_path to model.get_events_dataframe — text is automatically converted to speech and transcribed to obtain word-level timings.

For a full walkthrough with brain visualizations, see the Colab demo notebook.

Installation

Basic (inference only):

pip install -e .

With brain visualization:

pip install -e ".[plotting]"

With training dependencies (PyTorch Lightning, W&B, etc.):

pip install -e ".[training]"

Mac reward-node setup (shared-repo scripts)

If you want to run TRIBE reward inference on Mac and keep fine-tuning on a separate GPU machine, use the helper scripts in scripts/:

# from repo root
bash scripts/setup_mac_tribe.sh

# activate and load stable env vars
source .venv_mac_tribe/bin/activate
source .venv_mac_tribe/mac_tribe_env.sh

Connection and access tests

Quick connectivity/auth checks (Python, network, HF repo access):

python scripts/test_mac_tribe_setup.py --mode quick

Full TRIBE smoke test (loads model and runs one prediction):

python scripts/test_mac_tribe_setup.py --mode full --device cpu

Optional MPS test after CPU success:

python scripts/test_mac_tribe_setup.py --mode full --device mps

Online reward mode across Mac + Windows

Once the Mac environment is set up, you can use the Mac as a live TRIBE reward server and the Windows GPU machine as the policy-training node.

Start the reward server on the Mac:

source .venv_mac_tribe/bin/activate
source .venv_mac_tribe/mac_tribe_env.sh
python scripts/mac_tribe_reward_server.py --host 0.0.0.0 --port 8000 --device cpu

Check the server from Windows:

Invoke-WebRequest http://192.168.1.119:8000/health

Run online PPO on Windows using the Mac reward server:

python tribev2/unsloth_ppo_train.py --reward-url http://192.168.1.119:8000 --steps 10 --batch-size 1 --max-new-tokens 16

In this mode, the Windows machine generates text and sends the exact generated response to the Mac TRIBE server for live reward scoring before each PPO update.

Training a model from scratch

1. Set environment variables

Configure data/output paths and Slurm partition (or edit tribev2/grids/defaults.py directly):

export DATAPATH="/path/to/studies"
export SAVEPATH="/path/to/output"
export SLURM_PARTITION="your_partition"

2. Authenticate with HuggingFace

The text encoder requires access to the gated LLaMA 3.2-3B model:

huggingface-cli login

Create a read access token and paste it when prompted.

3. 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

Open Mic preprocessing (variable-length joke windows)

To start the humor project with an Open-Mic-first pipeline, prepare sentence-level transcripts into short, medium, and long segments:

python -m tribev2.openmic_prep \
  --input /path/to/openmic_sentences.csv \
  --output /path/to/openmic_segments.parquet \
  --summary /path/to/openmic_segments_summary.json

Expected input columns:

  • comedian_id
  • set_id
  • sentence_id
  • text
  • start_sec
  • end_sec
  • one of: laughter_sec or laughter_score

What this command does:

  • builds variable-length segments (short: 1 sentence, medium: rolling 2-4, long: laughter-centered chunks)
  • computes a duration/token-normalized humor_target
  • assigns deterministic leakage-safe train/val/test splits by set_id
  • reports balance by split and length_bin

End-to-end local loop (HF Open Mic text subset)

# 1) Convert HF Open Mic text files into sentence-level rows
python -m tribev2.openmic_from_hf \
  --input-root data/openmic_hf \
  --output data/openmic_hf/openmic_sentences.csv

# 2) Build variable-length segments and deterministic splits
python -m tribev2.openmic_prep \
  --input data/openmic_hf/openmic_sentences.csv \
  --output data/openmic_hf/openmic_segments.parquet \
  --summary data/openmic_hf/openmic_segments_summary.json

# 3) Train a lightweight local baseline on prepared segments
python -m tribev2.openmic_train \
  --data data/openmic_hf/openmic_segments.parquet \
  --out data/openmic_hf/runs/baseline

# 4) Build reward table (phase 2)
python -m tribev2.openmic_reward \
  --segments data/openmic_hf/openmic_segments.parquet \
  --model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --output data/openmic_hf/openmic_rewards.parquet \
  --summary data/openmic_hf/openmic_rewards_summary.json

# 5) Calibrate reward quality (phase 3)
python -m tribev2.openmic_calibrate \
  --input data/openmic_hf/openmic_rewards.parquet \
  --output data/openmic_hf/runs/calibration/report.json

# 6) Preference tuning with a lightweight DPO-like objective (phase 4)
python -m tribev2.openmic_dpo_lite \
  --reward-table data/openmic_hf/openmic_rewards.parquet \
  --out data/openmic_hf/runs/dpo_lite

# 7) Final baseline-vs-tuned evaluation and go/no-go decision (phase 5)
python -m tribev2.openmic_eval \
  --reward-table data/openmic_hf/openmic_rewards.parquet \
  --baseline-model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --tuned-model data/openmic_hf/runs/dpo_lite/openmic_dpo_lite_model.pt \
  --output data/openmic_hf/runs/final_eval/report.json

# 8) Build blind A/B test sets for human humor preference (next step)
python -m tribev2.openmic_ab_test generate \
  --data data/openmic_hf/openmic_segments.parquet \
  --baseline-model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --tuned-model data/openmic_hf/runs/dpo_lite/openmic_dpo_lite_model.pt \
  --generator-model distilgpt2 \
  --n-prompts 30 \
  --candidates-per-prompt 8 \
  --out data/openmic_hf/runs/ab_test

# 9) After filling votes (A/B/TIE) in ab_blind.jsonl, score outcomes
python -m tribev2.openmic_ab_test score \
  --blind-ab data/openmic_hf/runs/ab_test/ab_blind.jsonl \
  --answer-key data/openmic_hf/runs/ab_test/ab_answer_key.jsonl

This baseline is intentionally lightweight and local-friendly; it validates the full data/training/evaluation loop before larger reward-model and RL stages.

Notes on humor testing:

  • The current pipeline trains humor scorers/rankers; it does not yet fine-tune a generator model.
  • The A/B harness tests whether tuned scoring selects better candidates from the same generated candidate pool.
  • Use human preference votes as the final decision metric for "is it funnier?".

Project structure

tribev2/
├── main.py              # Experiment pipeline: Data, TribeExperiment
├── model.py             # FmriEncoder: Transformer-based multimodal→fMRI model
├── pl_module.py         # PyTorch Lightning training module
├── demo_utils.py        # TribeModel and helpers for inference from text/audio/video
├── eventstransforms.py  # Custom event transforms (word extraction, chunking, …)
├── utils.py             # Multi-study loading, splitting, subject weighting
├── utils_fmri.py        # Surface projection (MNI / fsaverage) and ROI analysis
├── grids/
│   ├── defaults.py      # Full default experiment configuration
│   └── test_run.py      # Quick local test entry point
├── plotting/            # Brain visualization (PyVista & Nilearn backends)
└── studies/             # Dataset definitions (Algonauts2025, Lahner2024, …)

Contributing to open science

If you use this software, please share your results with the broader research community using the following citation:

@article{dAscoli2026TribeV2,
  title={A foundation model of vision, audition, and language for in-silico neuroscience},
  author={d'Ascoli, St{\'e}phane and Rapin, J{\'e}r{\'e}my and Benchetrit, Yohann and Brookes, Teon and Begany, Katelyn and Raugel, Jos{\'e}phine and Banville, Hubert and King, Jean-R{\'e}mi},
  year={2026}
}

License

This project is licensed under CC-BY-NC-4.0. See LICENSE for details.

Contributing

See CONTRIBUTING.md for how to get involved.

yautungsky/braindance

Python

0

5 commits

updated Apr 12, 2026

See the code

README

TRIBE v2

A Foundation Model of Vision, Audition, and Language for In-Silico Neuroscience

Open In Colab License: CC BY-NC 4.0 Python 3.10+

📄 Paper ▶️ Demo | 🤗 Weights

TRIBE v2 is a deep multimodal brain encoding model that predicts fMRI brain responses to naturalistic stimuli (video, audio, text). It combines state-of-the-art feature extractors — LLaMA 3.2 (text), V-JEPA2 (video), and Wav2Vec-BERT (audio) — into a unified Transformer architecture that maps multimodal representations onto the cortical surface.

Quick start

Load a pretrained model from HuggingFace and predict brain responses to a video:

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)

Predictions are for the "average" subject (see paper for details) and live on the fsaverage5 cortical mesh (~20k vertices). You can also pass text_path or audio_path to model.get_events_dataframe — text is automatically converted to speech and transcribed to obtain word-level timings.

For a full walkthrough with brain visualizations, see the Colab demo notebook.

Installation

Basic (inference only):

pip install -e .

With brain visualization:

pip install -e ".[plotting]"

With training dependencies (PyTorch Lightning, W&B, etc.):

pip install -e ".[training]"

Mac reward-node setup (shared-repo scripts)

If you want to run TRIBE reward inference on Mac and keep fine-tuning on a separate GPU machine, use the helper scripts in scripts/:

# from repo root
bash scripts/setup_mac_tribe.sh

# activate and load stable env vars
source .venv_mac_tribe/bin/activate
source .venv_mac_tribe/mac_tribe_env.sh

Connection and access tests

Quick connectivity/auth checks (Python, network, HF repo access):

python scripts/test_mac_tribe_setup.py --mode quick

Full TRIBE smoke test (loads model and runs one prediction):

python scripts/test_mac_tribe_setup.py --mode full --device cpu

Optional MPS test after CPU success:

python scripts/test_mac_tribe_setup.py --mode full --device mps

Online reward mode across Mac + Windows

Once the Mac environment is set up, you can use the Mac as a live TRIBE reward server and the Windows GPU machine as the policy-training node.

Start the reward server on the Mac:

source .venv_mac_tribe/bin/activate
source .venv_mac_tribe/mac_tribe_env.sh
python scripts/mac_tribe_reward_server.py --host 0.0.0.0 --port 8000 --device cpu

Check the server from Windows:

Invoke-WebRequest http://192.168.1.119:8000/health

Run online PPO on Windows using the Mac reward server:

python tribev2/unsloth_ppo_train.py --reward-url http://192.168.1.119:8000 --steps 10 --batch-size 1 --max-new-tokens 16

In this mode, the Windows machine generates text and sends the exact generated response to the Mac TRIBE server for live reward scoring before each PPO update.

Training a model from scratch

1. Set environment variables

Configure data/output paths and Slurm partition (or edit tribev2/grids/defaults.py directly):

export DATAPATH="/path/to/studies"
export SAVEPATH="/path/to/output"
export SLURM_PARTITION="your_partition"

2. Authenticate with HuggingFace

The text encoder requires access to the gated LLaMA 3.2-3B model:

huggingface-cli login

Create a read access token and paste it when prompted.

3. 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

Open Mic preprocessing (variable-length joke windows)

To start the humor project with an Open-Mic-first pipeline, prepare sentence-level transcripts into short, medium, and long segments:

python -m tribev2.openmic_prep \
  --input /path/to/openmic_sentences.csv \
  --output /path/to/openmic_segments.parquet \
  --summary /path/to/openmic_segments_summary.json

Expected input columns:

  • comedian_id
  • set_id
  • sentence_id
  • text
  • start_sec
  • end_sec
  • one of: laughter_sec or laughter_score

What this command does:

  • builds variable-length segments (short: 1 sentence, medium: rolling 2-4, long: laughter-centered chunks)
  • computes a duration/token-normalized humor_target
  • assigns deterministic leakage-safe train/val/test splits by set_id
  • reports balance by split and length_bin

End-to-end local loop (HF Open Mic text subset)

# 1) Convert HF Open Mic text files into sentence-level rows
python -m tribev2.openmic_from_hf \
  --input-root data/openmic_hf \
  --output data/openmic_hf/openmic_sentences.csv

# 2) Build variable-length segments and deterministic splits
python -m tribev2.openmic_prep \
  --input data/openmic_hf/openmic_sentences.csv \
  --output data/openmic_hf/openmic_segments.parquet \
  --summary data/openmic_hf/openmic_segments_summary.json

# 3) Train a lightweight local baseline on prepared segments
python -m tribev2.openmic_train \
  --data data/openmic_hf/openmic_segments.parquet \
  --out data/openmic_hf/runs/baseline

# 4) Build reward table (phase 2)
python -m tribev2.openmic_reward \
  --segments data/openmic_hf/openmic_segments.parquet \
  --model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --output data/openmic_hf/openmic_rewards.parquet \
  --summary data/openmic_hf/openmic_rewards_summary.json

# 5) Calibrate reward quality (phase 3)
python -m tribev2.openmic_calibrate \
  --input data/openmic_hf/openmic_rewards.parquet \
  --output data/openmic_hf/runs/calibration/report.json

# 6) Preference tuning with a lightweight DPO-like objective (phase 4)
python -m tribev2.openmic_dpo_lite \
  --reward-table data/openmic_hf/openmic_rewards.parquet \
  --out data/openmic_hf/runs/dpo_lite

# 7) Final baseline-vs-tuned evaluation and go/no-go decision (phase 5)
python -m tribev2.openmic_eval \
  --reward-table data/openmic_hf/openmic_rewards.parquet \
  --baseline-model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --tuned-model data/openmic_hf/runs/dpo_lite/openmic_dpo_lite_model.pt \
  --output data/openmic_hf/runs/final_eval/report.json

# 8) Build blind A/B test sets for human humor preference (next step)
python -m tribev2.openmic_ab_test generate \
  --data data/openmic_hf/openmic_segments.parquet \
  --baseline-model data/openmic_hf/runs/baseline/openmic_baseline_model.pt \
  --tuned-model data/openmic_hf/runs/dpo_lite/openmic_dpo_lite_model.pt \
  --generator-model distilgpt2 \
  --n-prompts 30 \
  --candidates-per-prompt 8 \
  --out data/openmic_hf/runs/ab_test

# 9) After filling votes (A/B/TIE) in ab_blind.jsonl, score outcomes
python -m tribev2.openmic_ab_test score \
  --blind-ab data/openmic_hf/runs/ab_test/ab_blind.jsonl \
  --answer-key data/openmic_hf/runs/ab_test/ab_answer_key.jsonl

This baseline is intentionally lightweight and local-friendly; it validates the full data/training/evaluation loop before larger reward-model and RL stages.

Notes on humor testing:

  • The current pipeline trains humor scorers/rankers; it does not yet fine-tune a generator model.
  • The A/B harness tests whether tuned scoring selects better candidates from the same generated candidate pool.
  • Use human preference votes as the final decision metric for "is it funnier?".

Project structure

tribev2/
├── main.py              # Experiment pipeline: Data, TribeExperiment
├── model.py             # FmriEncoder: Transformer-based multimodal→fMRI model
├── pl_module.py         # PyTorch Lightning training module
├── demo_utils.py        # TribeModel and helpers for inference from text/audio/video
├── eventstransforms.py  # Custom event transforms (word extraction, chunking, …)
├── utils.py             # Multi-study loading, splitting, subject weighting
├── utils_fmri.py        # Surface projection (MNI / fsaverage) and ROI analysis
├── grids/
│   ├── defaults.py      # Full default experiment configuration
│   └── test_run.py      # Quick local test entry point
├── plotting/            # Brain visualization (PyVista & Nilearn backends)
└── studies/             # Dataset definitions (Algonauts2025, Lahner2024, …)

Contributing to open science

If you use this software, please share your results with the broader research community using the following citation:

@article{dAscoli2026TribeV2,
  title={A foundation model of vision, audition, and language for in-silico neuroscience},
  author={d'Ascoli, St{\'e}phane and Rapin, J{\'e}r{\'e}my and Benchetrit, Yohann and Brookes, Teon and Begany, Katelyn and Raugel, Jos{\'e}phine and Banville, Hubert and King, Jean-R{\'e}mi},
  year={2026}
}

License

This project is licensed under CC-BY-NC-4.0. See LICENSE for details.

Contributing

See CONTRIBUTING.md for how to get involved.