README.md

---
license: apache-2.0
---
# TEDDY: A Family of Foundation Models for Single-Cell Biology

This repository provides open-source code and configurations supporting the **TEDDY** project, as described in:

> **[TEDDY: A FAMILY OF FOUNDATION MODELS FOR UNDERSTANDING SINGLE CELL BIOLOGY](https://arxiv.org/abs/2503.03485)**

TEDDY leverages large-scale single-cell RNA sequencing (scRNA-seq) data (~116 million cells) to train transformer-based models. These models capture disease-related signals and generalize to diverse downstream tasks, including cross-donor and cross-disease classification.

---

## Table of Contents

1. [Introduction](#introduction)
2. [Project Goals & Paper Summary](#project-goals--paper-summary)
3. [Pipeline Overview](#pipeline-overview)
4. [Installation](#installation--setup)
5. [Detailed Steps](#detailed-steps)
   1. [Preprocessing & Tokenization](#1-preprocessing--tokenization)
   2. [Loading TEDDY Models](#2-loading-teddy-models)
6. [Running sample scripts on sample data](#running-sample-scripts-on-sample-data)
7. [Running unit tests with pytest](#running-unit-tests-with-pytest)
8. [Reference](#reference)

---

## Introduction

Single-cell RNA sequencing data can span hundreds of millions of cells, each expressing thousands of genes. **TEDDY** (*T*ransformer for *E*nabling *D*rug *D*iscovery) adapts masked language modeling and ontology classification to gene expression. By scaling both **data volume** and **model capacity** (up to ~400M parameters), TEDDY learns robust biological features that generalize to **unseen diseases**, **unseen donors**, and more.

---

## Project Goals & Paper Summary

Refer to the [paper](https://arxiv.org/abs/2503.03485) for full technical details. Key highlights:

- **Data**: 116 million single cells, spanning multiple tissues, diseases, and human/mouse species.
- **Models**:
  - **TEDDY-G** (rank-based encoding)
  - **TEDDY-X** (binned encoding)
  - Sizes range from 10M to 400M parameters.
- **Annotation Supervision**: Additional labels (disease, tissue, cell type, etc.) further refine model representations.
- **Benchmarks**: “Held-out donors” and “held-out diseases” classification tasks showed significant gains over alternative foundation models.

(Note: This release only includes the most performant models: TEDDY-G 70M, TEDDY-G 160M, and TEDDY-G 400M)
---

## Pipeline Overview

**TEDDY** pipeline involves three steps:

1. **Preprocessing**
   - Load `.h5ad` files, remove low-quality cells, normalize expression counts to 10000,
   and median normalize.
   - Outputs a “processed” `.h5ad` file.

2. **Tokenization**
   - Converts each cell’s expression profile into integer tokens or rank-based embeddings.
   - Can embed metadata tokens that can be used as ontologies in the model (e.g., `<disease>`, `<tissue_type>`, `<sex>`, `<cell_type>`) if needed.

3. **Model Inference and Training**
   - Uses the tokenized dataset to generate embeddigns for cells and genes.
   - Uses the tokenized dataset for masked language modeling plus ontology classification.
   - Model config examples live in dedicated config files for relevant architectures.

---

## Installation & Setup

**Building Your Environment**

### 1. Clone the Repository
First, clone the repository to your local machine:
```bash
git clone https://huggingface.co/Merck/TEDDY
cd TEDDY
```

### 2. Environment Setup

- Fine-tuning and pretraining of these models were conducted on GPUs, so ensure your instance is properly configured before working with large datasets.

- Ensure you have ***Python 3.11.10*** installed. You can use `pyenv` to manage Python versions:
  ```bash
  pyenv install 3.11.10
  pyenv local 3.11.10
  ```
  - More details on how to use `pyenv`: [pyenv documentation](https://github.com/pyenv/pyenv)


- If you don’t already have ***Poetry*** installed, you can install it using the following command:
  ```bash
  curl -sSL https://install.python-poetry.org | python3 -
  export PATH="/PATH/TO/YOUR/USER/.local/bin:$PATH"
  ```
- Check that poetry uses the correct python version:
  ```bash
  pyenv which python
  ```
  - Change to correct version by running:
    ```bash
    poetry env use /PATH/TO/YOUR/USER/.pyenv/versions/3.11.10/bin/python
    ```
- Run the following command to build the project and install its dependencies:
  ```bash
  poetry build
  poetry install
  ```
- Once the setup is complete, you can use the package.

---

## Detailed Steps

There are three ways to run **Preprocessing** and **Tokenization**:
1. **Directly in Python** (importing the scripts)
2. **Command-Line Arguments** (using flags)
3. **JSON Config Files** (loading a `.json` with your parameters)

### 1. Preprocessing & Tokenization
#### Directly in Python

Detailed [README.md for Preprocessing](teddy/data_processing/preprocessing/README.md) and [README.md for Tokenization](teddy/data_processing/tokenization/README.md) can be found in the related module folders.

**Preprocessing example**:

```
from teddy.data_processing.preprocessing.preprocess import preprocess

preprocessing_config = {
  "min_gene_counts": 225,
  "remove_assays": ["10x5' v1", "10x3' v1"],
  "max_mitochondrial_prop": 10,
  "remove_cell_types": [],
  "hvg_method": None,
  "normalized_total": 10000,
  "median_dict": "teddy/data_processing/utils/medians/data/teddy_gene_medians.json",
  "log1p": False,
  "compute_medians": False,
  "median_column": "index",
  "reference_id_only": False,
  "load_dir": "<PATH_TO_RAW_DATA_PARENT>",
  "save_dir": "<PATH_TO_PROCESSED_DATA_PARENT>",
}

preprocess(
  data_path="data/RAW_SAMPLES/my_data.h5ad",
  metadata_path="data/RAW_SAMPLES/my_data_metadata.json",
  hyperparameters=preprocessing_config
)
```

The above preprocessing arguments were used to preprocess the corpus used for pretraining TEDDY models.

**Tokenization example**:

```
from teddy.data_processing.tokenization.tokenization import tokenize

tokenizer_config = {
  "tokenizer_name_or_path": "teddy/models/teddy_g/400M",
  "gene_id_column": "index",
  "bio_annotations": True,
  "disease_mapping": "teddy/data_processing/utils/bio_annotations/data/mappings/all_filtered_disease_mapping.json",
  "tissue_mapping": "teddy/data_processing/utils/bio_annotations/data/mappings/all_filtered_tissue_mapping.json",
  "cell_mapping": "teddy/data_processing/utils/bio_annotations/data/mappings/all_filtered_cell_mapping.json",
  "sex_mapping": "teddy/data_processing/utils/bio_annotations/data/mappings/all_filtered_sex_mapping.json",
  "max_shard_samples": 500,
  "max_seq_len": 2048,
  "pad_length": 2048,
  "add_cls": False,
  "bins": 0,
  "continuous_rank": True,
  "truncation_method": "max",
  "add_disease_annotation": False,
  "include_zero_genes": False,
  "load_dir": "<PATH_TO_PROCESSED_DATA_PARENT>",
  "save_dir": "<PATH_TO_TOKENIZED_DATA>"
}

tokenize(
  data_path="outputs/preprocessed/my_data_preprocessed.h5ad",
  metadata_path="outputs/preprocessed/my_data_preprocessed_metadata.json",
  hyperparameters=tokenizer_config
)
```

Above tokenization arguments were used for the Teddy models.

####  By Saving a `config.json` and Running It with Bash

Example **preprocess_config.json**:

```
{
  "min_gene_counts": null,
  "remove_assays": [],
  "max_mitochondrial_prop": null,
  "remove_cell_types": [],
  "hvg_method": null,
  "normalized_total": 10000,
  "median_dict": "teddy/data_processing/utils/medians/data/teddy_gene_medians.json",
  "log1p": false,
  "compute_medians": false,
  "median_column": "index",
  "reference_id_only": false,
  "load_dir": "<PATH_TO_RAW_DATA_PARENT>",
  "save_dir": "<PATH_TO_PROCESSED_DATA_PARENT>"
}
```
Run:
```
python teddy/data/preprocessing/preprocess.py \
  --data_path data/RAW_SAMPLES/my_data.h5ad \
  --metadata_path data/RAW_SAMPLES/my_data_metadata.json \
  --config preprocess_config.json
```
(Same idea for tokenization, e.g. tokenize_config.json, then `--config tokenize_config.json`.)

#### By Creating a `.sh` File and Executing It (With Poetry)

You can find an example in **scripts/preprocess_sample_data.sh**:
```
#!/bin/bash -l

# (Optional) Activate your Poetry environment
poetry shell

# 1) Generate a JSON config file on the fly
cat <<EOF > configs/my_preprocess_config.json
{
  "load_dir": "data",
  "save_dir": "data/processed",

  "min_gene_counts": null,
  "remove_assays": [],
  "max_mitochondrial_prop": null,
  "remove_cell_types": [],
  "hvg_method": null,
  "normalized_total": null,

  "median_dict": "teddy/data_processing/utils/medians/data/teddy_gene_medians.json",
  "log1p": false,
  "compute_medians": false,
  "median_column": "index",

  "reference_id_only": false
}
EOF

# 2) Call preprocess.py, explicitly passing data_path, metadata_path, and config_path
python teddy/data_processing/preprocessing/preprocess.py \
  --data_path data/sample_data.h5ad \
  --metadata_path data/sample_data_metadata.json \
  --config_path my_preprocess_config.json
```
Then do:
```
chmod +x preprocess_sample_data.sh
./preprocess_sample_data.sh
```
---
You can override any parameter by specifying command-line arguments, editing the `.json`, or updating the Python dictionary.

(Same idea for tokenization, e.g. use example in **scripts/tokenize_sample_data.sh**)

### 2. Loading TEDDY Models

If you want to load a trained TEDDY model in your Python code, you can do so with the following snippet:

```
from teddy.models.model_directory import get_architecture, model_dict

model_name_or_path = 'teddy/models/teddy_g/400M' # or local path to model files
arch = get_architecture(model_name_or_path)
config_cls = model_dict[arch]["config_cls"]
model_cls = model_dict[arch]["model_cls"]

# Load the configuration and model
config = config_cls.from_pretrained(model_name_or_path)
model = model_cls.from_pretrained(model_name_or_path, config=config)
# model is now ready for inference or further fine-tuning
```
You can then perform inference, fine-tuning, or evaluation with the model object as needed.

## Running sample scripts on sample data

In the `scripts` directory of this repository, sample code has been included with which to preprocess and tokenize the sample data in the `data` directory. To switch this out for your own data, simply replace the data within the `data` directory with your data and rename file paths within the scripts as needed.

To run the scripts included, run the following commands from the root of the `teddy-models` repository.

```
chmod +x scripts/*
./scripts/preprocess_sample_data.sh
./scripts/tokenize_sample_data.sh
```

## Running unit tests with pytest

To run the unit tests in the repository, you can run `poetry run pytest`. The tests should all pass, but receiving runtime warnings is expected behavior with the simulated data for the tests.

## Reference

Reference to cite when you use TEDDY:

```
@misc{chevalier2025teddyfamilyfoundationmodels,
      title={TEDDY: A Family Of Foundation Models For Understanding Single Cell Biology},
      author={Alexis Chevalier and Soumya Ghosh and Urvi Awasthi and James Watkins and Julia Bieniewska and Nichita Mitrea and Olga Kotova and Kirill Shkura and Andrew Noble and Michael Steinbaugh and Julien Delile and Christoph Meier and Leonid Zhukov and Iya Khalil and Srayanta Mukherjee and Judith Mueller},
      year={2025},
      eprint={2503.03485},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2503.03485},
```