PyTorch implementation of scCello, a cell-ontology guided transcriptome foundation model (TFM) for single cell RNA-seq data. Authored by Xinyu Yuan, and Zhihao Zhan.
scCello enhances transcriptome foundation models (TFMs) by integrating cell ontology graphs into pre-training, addressing the limitation of treating cells as independent entities. By incorporating cell-level objectives: cell-type coherence loss and ontology alignment loss, scCello demonstrate superior or competitive generalization and transferability capability over the existing TFMs on biologically important tasks including identifying novel cell types of unseen cells, prediction of cell-type-specific marker genes, and cancer drug responses.
This repository is based on PyTorch 2.0 and Python 3.9.
Table of contents:
- Cell-type Specific Learning: Utilizes cell-type coherence loss to learn specific gene expression patterns relevant to each cell type.
- Ontology-aware Modeling: Employs ontology alignment loss to understand and preserve the hierarchical relationships among different cell types.
- Large-scale Pre-training: Trained on over 22 million cells from the CellxGene database, ensuring robust and generalizable models.
- Advanced Generalization and Transferability: Demonstrates superior performance on various biologically significant tasks such as identifying novel cell types and predicting cell-type-specific marker genes.
- Feb 5th, 2025: scCello code released!
- Oct 1st, 2024: scCello got accepted at NeurIPS 2024!
- Aug 22nd, 2024: scCello preprint release on arxiv!
You may install the dependencies via the following bash command.
conda install pytorch==2.0.1 pytorch-cuda=11.7 -c pytorch -c nvidia pip install transformers[torch] pip install easydict pip install psutil pip install wandb pip install pytz pip install ipdb pip install pandas pip install datasets pip install torchmetrics pip install rdflib pip install hickle pip install anndata pip install scikit-learn pip install scanpy pip install scib conda install -c conda-forge cupy conda install rapidsai::cuml conda install -c rapidsai -c conda-forge -c nvidia cugraphQuick start guide to load scCello checkpoint:
- for zero-shot inference tasks
from sccello.src.model_prototype_contrastive import PrototypeContrastiveForMaskedLM model = PrototypeContrastiveForMaskedLM.from_pretrained("katarinayuan/scCello-zeroshot", output_hidden_states=True) - for linear probing tasks (see details in sccello/script/run_cell_type_classification.py)
from sccello.src.model_prototype_contrastive import PrototypeContrastiveForSequenceClassification model_kwargs = { "num_labels": NUM_LABELS, # number of labels for classification "total_logging_steps": training_cfg["logging_steps"] * training_args.gradient_accumulation_steps, } model = PrototypeContrastiveForSequenceClassification.from_pretrained("katarinayuan/scCello-zeroshot", **model_kwargs) For downstreams, in-distribution (ID) data
# Note that some datasets are extremely large, use the following command to change data caching directory. The default is "~/.cache/huggingface/datasets/". export HF_HOME="/path/to/another/directory/datasets" from sccello.src.utils import data_loading # pre-training data & D^{id} train_dataset = load_dataset("katarinayuan/scCello_pretrain_unsplitted")["train"] train_dataset, indist_test_data = train_dataset.train_test_split(test_size=0.001, seed=237) # seed used in scCello # D_1^{ct} & D_2^{ct} d1_ct, d2_ct = data_loading.get_fracdata("celltype", "frac100", False, False) # D_1^{ts} & D_2^{ts} d1_ts, d2_ts = data_loading.get_fracdata("tissue", "frac100", False, False) # D_1^{dn} & D_2^{dn} d1_dn, d2_dn = data_loading.get_fracdata("donor", "frac100", False, False) Example data for transforming h5ad format to huggingface format. For building pre-training datasets and downstream datasets, we downloaded a series of human h5ad data from CellxGene
pip install gdown cd ./data/example_h5ad/ gdown https://drive.google.com/uc?id=1UsbkhmZwSDWTgY4die60fHvzL_FnXtWEThe sccello/script folder contains all executable files.
General configurations:
pretrained_ckpt=katarinayuan/scCello-zeroshot output_dir=/home/xinyu402/single_cell_output/ wandb_run_name=test python ./sccello/script/run_data_transformation.py python ./sccello/script/run_cell_type_clustering.py --pretrained_ckpt $pretrained_ckpt --wandb_run_name $wandb_run_name --output_dir $output_dir # Linear Probing training_type=linear_probing # or Train from Scratch without Loading the Pre-trained Model # training_type=from_scratch_linear torchrun ./sccello/script/run_cell_type_classification.py --pretrained_ckpt $pretrained_ckpt --training_type $training_type --wandb_run_name $wandb_run_name --further_downsample 0.01 --output_dir $output_dir python ./sccello/script/run_novel_cell_type_classification.py --pretrained_ckpt $pretrained_ckpt --wandb_run_name $wandb_run_name --indist_repr_path ./embedding_storage/cellreprs_indist_frac_celltype_data1.pkl --output_dir $output_dir python ./sccello/script/run_marker_gene_prediction.py --pretrained_ckpt $pretrained_ckpt --wandb_run_name $wandb_run_name --output_dir $output_dir python ./sccello/script/run_cancer_drug_response.py --pretrained_ckpt $pretrained_ckpt --wandb_run_name $wandb_run_name python -m torch.distributed.run --nproc_per_node=1 ./sccello/script/run_pretrain_prototype_contrastive.py --wandb_run_name pretrain_test If you find this codebase useful in your research, please cite the original papers.
The main scCello paper:
@inproceedings{yuancell, title={Cell ontology guided transcriptome foundation model}, author={Yuan, Xinyu and Zhan, Zhihao and Zhang, Zuobai and Zhou, Manqi and Zhao, Jianan and Han, Boyu and Li, Yue and Tang, Jian}, booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems} }