uce_hf

Universal Cell Embeddings

This repo is an UNOFFICIAL hugging implementation of the UCE (https://github.com/snap-stanford/UCE) model, designed for generate embeddings for single cell dataset. Based on the official repo of UCE, this repo provides scripts for pretraining, continue-training, and inference.

Installation

pip install -r requirements.txt

Dataset

A typical dataset folder containing .h5ad files looks like below:

├── dataset_uce     # dataset root
│   ├── human 		# species
│   	├── human-sample1.h5ad
│   	├── human-sample2.h5ad
│   	├── ...
│   ├── mouse
│   	├── mouse-sample1.h5ad
│   	├── mouse-sample2.h5ad
│   	├── ...
│   ├── ...

Samples of the same species belong to the same species folder. The filenames of the species folder are used to map to the protein embedding model. See preprocessors.py for details.

More test samples can be found from the official dataset source

Train

We provide a simple example to show how to continue-train a 4 layer uce model. To train a model from scratch, one can simply skip the model loading part.

python train_hf.py --dataset_folder dataset/ \
                --output_dir ./results/uce_ckpts/test-1 \ 
                --dir test_output/ \
                --model_loc model_files/4layer_model.torch \
                --pretrained_model model-4 \
                --logging_dir logs/exp1 \
                --batch_size 4

Also, we can start the pretrain using deepspeed

deepspeed train_hf.py --dataset_folder dataset/ \
                --output_dir ./results/test1 \
                --dir test_output/ \
                --model_loc model_files/4layer_model.torch \
                --pretrained_model model-4 \
                --logging_dir /logs \
                --batch_size 4 \
                --deepspeed_config ds_configs/zero2.json