Tabula

Preprint

Predictive single cell foundation model for gene regulation and aging with privacy-preserving tabular learning

Jiayuan Ding, Jianhui Lin, Ziyang Miao, Nils Mechtel, Shiyu Jiang, Yixin Wang, Zhaoyu Fang, Jorge D. Martin-Rufino, Chen Weng, Reuben Saunders, Weize Xu, Jonathan S. Weissman, Min Li, Jiliang Tang, Wei Ouyang, Yuancheng Ryan Lu, Xiaojie Qiu

Overview

A privacy-preserving predictive foundation model for single-cell transcriptomics, leveraging federated learning and tabular modeling.

Project Structure

Tabula/  
├── resource/                      
│   ├── dataset/                   # Processed pretrian datasets  
│   ├── finetune_framework_x.yaml  # The configuration of downstream task  
│   ├── vocab.json                 # Genetic vocabulary

├── tabula/                      
│   ├── downstream/            # Downstream task implementations
│   ├── model/                 
│   │   ├── encoding/          # Single-cell data embedding  
│   │   ├── transformer/       # Transformer backbone  
│   ├── loss.py                # Training loss  

│   ├── training/              # Pre-training  
│   │   ├── config.py          # Configuration  
│   │   ├── data_loader.py     # Multi-client data loader  
│   │   ├── federater.py       # Federated framework  
│   │   ├── pretrainer.py      # PyTorch Lightning training framework  

├── tests/                     # Unit tests

├── tutorials/                 # Usage examples for downstream task

├── requirements.txt           # Python dependencies  
├── README.md                  # Project description file
└── LICENSE  

Installation

• CUDA >= 11.7
• Python >= 3.9
• flash-attn >= 2.3.5
• mpi4py >= 3.1.4
• Required dependencies are listed in requirements.txt

Clone the repository:

$ git clone this-repo-url
$ cd tabula

Create your conda conda environment:

$ conda install -n tabula python=3.9

Install the torch:

$ pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117

Install dependencies:

$ pip install -r requirements.txt

To install mpi4py, follow these steps:

$ conda install mpi4py==3.1.4

To install flash-attention2, follow these steps, (For more information, check out flash-attention):

$ MAX_JOBS=4 pip install flash-attn==2.3.5 --no-build-isolation

Quick Start Tutorial

Please see our example code in tutorials.

Tabula Development Process

• Follow feature-staging-main review process
  • create a specific branch for new feature
  • implement and test on your branch; add unit tests
  • create pull request
  • discuss with lab members and merge into the main branch once all checks pass
• Follow python Google code style

Code quality

• File and function docstrings should be written in Google style
• We use black to automatically format code in a standardized format. To ensure that any code changes are up to standard, use pre-commit as such.

# Run the following two lines ONCE.
$ pip install pre-commit
$ pre-commit install
$ pre-commit run --all-files

Then, all future commits will call black automatically to format the code. Any code that does not follow the standard will cause a check to fail.

Contact

For questions, feedback, or collaboration opportunities, please contact Xiaojie Qiu at xiaojie@stanford.edu and Jiayuan Ding at jiayuandingjsnu@gmail.com.