Metadata-Version: 2.1
Name: tahoex
Version: 0.1.2
Summary: Tahoe-x1: Perturbation trained virtual cell models from Tahoe Therapeutics
Author: Vevo Therapeutics
License: Proprietary
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: awscli>=1.32
Requires-Dist: llm-foundry[gpu]==0.17.1
Requires-Dist: mosaicml-streaming
Requires-Dist: ninja
Requires-Dist: cmake
Requires-Dist: packaging
Requires-Dist: torch==2.5.*
Requires-Dist: llvmlite
Requires-Dist: numba
Requires-Dist: natsort
Requires-Dist: scanpy
Requires-Dist: pynndescent
Requires-Dist: umap-learn
Requires-Dist: anndata
Requires-Dist: h5py
Provides-Extra: all
Requires-Dist: pre-commit<4,>=3.4.0; extra == "all"
Requires-Dist: ipykernel; extra == "all"
Requires-Dist: transformer-engine@ git+https://github.com/NVIDIA/TransformerEngine.git@stable ; extra == "all"
Requires-Dist: packaging<23,>=21; extra == "all"
Requires-Dist: toml<0.11,>=0.10.2; extra == "all"
Requires-Dist: tahoex[dev]; extra == "all"
Provides-Extra: dev
Requires-Dist: pre-commit<4,>=3.4.0; extra == "dev"
Requires-Dist: toml<0.11,>=0.10.2; extra == "dev"
Requires-Dist: ipykernel; extra == "dev"
Requires-Dist: packaging<23,>=21; extra == "dev"
Provides-Extra: gpu
Requires-Dist: transformer-engine@ git+https://github.com/NVIDIA/TransformerEngine.git@stable ; extra == "gpu"

<!-- SETUPTOOLS_LONG_DESCRIPTION_HIDE_BEGIN -->
<p align="center">
  <a href="https://github.com/vevotx/mosaicfm">
    <picture>
      <img alt="tahoe-therapeutics" src="./assets/tahoe-navy-logo.png" width="95%">
    </picture>
  </a>
</p>
<!-- SETUPTOOLS_LONG_DESCRIPTION_HIDE_END -->

<p align="center">
<a href="https://github.com/astral-sh/ruff"><img alt="Linter: Ruff" src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json"></a>
<a href="https://github.com/vevotx/mosaicfm/blob/main/LICENSE">
        <img alt="License" src="https://img.shields.io/badge/License-Apache%202.0-green.svg">
    </a>
<a href="https://github.com/psf/black"><img alt="Code style: black" src="https://img.shields.io/badge/code%20style-black-000000.svg"></a>
</p>
<br />

# TahoeX

This is the internal codebase for the **TahoeX** series of single-cell RNA-seq foundation models 
developed by Tahoe Therapeutics. Our repository follows a similar structure to [llm-foundry](https://github.com/mosaicml/llm-foundry/tree/main) 
and imports several utility functions from it. Please follow the developer guidelines if you are 
contributing to this repository. For main results and documentation, please refer to the results section. 
If you are looking to train or finetune a model on single-cell data, please refer to the training section.

The repository is organized as follows:
* `tahoex/` contains the building blocks for the TahoeX models.
  * `tahoex/model/blocks` Building block modules that may be used across models
  * `tahoex/model/model` Full architectures subclassed from [ComposerModel](https://docs.mosaicml.com/projects/composer/en/latest/composer_model.html)
  * `tahoex/tasks/` Helper functions to use in downstream applications, such as embedding extraction.
  * `tahoex/tokenizer` Vocabulary building and tokenization functions.
  * `tahoex/data` Data loaders and collators
  * `tahoex/utils` Miscellaneous utility functions such to dowload files from s3 etc.
* `scripts/` contains scripts to train/evaluate models and to build datasets.
  * `scripts/train.py` Script to train a model. Accepts a yaml file or command line arguments for specifying job parameters.
  * `scripts/prepare_for_inference.py` Script to save a model for inference by packaging it with the vocabulary and saving metadata.
  * `scripts/depmap` Scripts to run the depmap benchmark.
* `mcli` yaml files to configure and launch runs on the MosaicML platform.
* `runai` yaml files to configure and launch runs on RunAI.
* `tutorials` Notebooks to demonstrate some applications of the models.

## Hardware and Software Requirements

We have tested our code on NVIDIA A100 and H100 GPUs with CUDA 12.1. 
At the moment, we are also restricted to use a version of llm-foundry no later v0.6.0, since support for the triton 
implementation of flash-attention was removed in [v0.7.0](https://github.com/mosaicml/llm-foundry/releases/tag/v0.7.0).

We support launching runs on the MosaicML platform as well as on local machines through RunAI.
The recommended method for using TahoeX is to use the pre-built [vevotx/mosaicfm](https://hub.docker.com/repository/docker/vevotx/mosaicfm/) docker image.

Currently, we have the following images available:

| Image Name                     | Base Image | Description         |
|--------------------------------|-------------|-------------------------------------------------|
| [`vevotx/mosaicfm:1.1.0`](https://github.com/vevotx/vevo-docker/tree/main/ml_docker_vevo_mosaicfm) | docker.io/mosaicml/llm-foundry:2.2.1_cu121_flash2-813d596 | Image used for TahoeX (current release) |

## Installation

### With docker
```shell
git clone https://github.com/tahoebio/tahoe-x1.git
cd tahoe-x1
pip install -e .
```
### Without docker
```shell
git clone https://github.com/tahoebio/tahoe-x1.git
cd tahoe-x1
mamba env create -f envs/mosaicfm_env.yml
mamba activate tahoe-x1
pip install -e . --no-deps # Inside the tahoex directory
```
> [!NOTE]  
> If you are on an H100 GPU you may see `'sm_90' is not a recognized processor for this target (ignoring processor)`. This is expected and safe to ignore.


## Datasets

The following datasets are used for training and evaluation:

| Dataset Path                             | Description                                                                                                     |
|------------------------------------------|-----------------------------------------------------------------------------------------------------------------|
| `s3://vevo-ml-datasets/vevo-scgpt/datasets/cellxgene_primary_2024-04-29_MDS/`     | MDS dataset comprising ~45M cells from Apr 2024 release by CellxGene and Vevo dataset 35 (resistance-is-futile) |
| `s3://vevo-ml-datasets/vevo-scgpt/datasets/cellxgene_primary_2023-12-15_MDS_v2/` | MDS dataset comprising ~34M cells from Dec 2023 release by CellxGene.                                           |
 |`s3://vevo-ml-datasets/umair/scgpt-depmap/`                        | Root folder containing Depmap dataset and model predictions                                                     |
|`s3://vevo-drives/drive_3/ANALYSIS/analysis_107/`                  | Root folder MSigDB data and model predictions                                                                   |

## Pre-trained Models

| Model Name                  | Run Name                           | Path to Checkpoints                                           | WandB id |
|-----------------------------|------------------------------------|---------------------------------------------------------------|----------|
| **TX1-1.3B**           | scgpt-1_3b-2048-prod               | `s3:/vevo-scgpt/models/scgpt-1_3b-2048-prod/`                 | lv6jl8kl |
| **TX1-70M**            | scgpt-70m-1024-fix-norm-apr24-data | `s3:/vevo-scgpt/models/scgpt-70m-1024-fix-norm-apr24-data/`    | 55n5wvdm |
| **TX1-25M**            | scgpt-25m-1024-fix-norm-apr24-data | `s3:/vevo-scgpt/models/scgpt-25m-1024-fix-norm-apr24-data/`    | bt4a1luo |
| **TX1-9M**             | sscgpt-test-9m-full-data           | `s3:/vevo-scgpt/models/scgpt-test-9m-full-data/`              | di7kyyf1 |

## Results
Links to evaluations and benchmarks are provided below:
 - [Depmap](scripts/depmap/README.md)
 - [MSigDB](https://github.com/vevotx/shreshth_sandbox/tree/main/analysis/04_msigdb_benchmark/README.md)

Please refer to our  technical report  for detailed results and analysis: [Internal Link](https://drive.google.com/drive/u/1/folders/1KeAXZ9zNYh4uHbLL5XUMmreAkHXW4yXo)

## Developer Guidelines
We use the black code style and the Ruff linter to mantain consistency across contributions.
Please set-up `pre-commit` and run the repository level hooks before committing any changes.
Please do not push to master directly. Create a new branch and open a pull request for review.
To set up pre-commit hooks, run the following command:
```shell
pip install pre-commit
pre-commit install
pre-commit run --all-files # Before committing
```
We also encourage new contributions to use type annotations and docstrings for functions and classes. In the future we
will add `pyright` and `pydocstyle` checks to the pre-commit hooks. We encourage the use of Google style docstrings.

If you will be launching any training/evaluation runs, please also make sure you have access to `s3`, `wandb` 
and `mcli`/`runai` by reaching out on #infrastructure.

## Acknowledgements
We would like to thank the developers of the following open-source projects:
 - [scGPT](https://github.com/bowang-lab/scGPT/tree/main)
 - [llm-foundry](https://github.com/mosaicml/llm-foundry)
 - [streaming](https://github.com/mosaicml/streaming)
 - [datasets](https://github.com/huggingface/datasets)
