README.md

metadata

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

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
2. Project Goals & Paper Summary
3. Pipeline Overview
4. Installation
5. Detailed Steps
   1. Preprocessing & Tokenization
   2. Loading TEDDY Models
6. Running sample scripts on sample data
7. Running unit tests with pytest
8. Reference

---

Introduction

Single-cell RNA sequencing data can span hundreds of millions of cells, each expressing thousands of genes. TEDDY (Transformer for Enabling Drug Discovery) 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 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:

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:

  pyenv install 3.11.10
  pyenv local 3.11.10
  

  • More details on how to use pyenv: pyenv documentation

• If you don’t already have Poetry installed, you can install it using the following command:

  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:

  pyenv which python
  

  • Change to correct version by running:

    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:

  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 and README.md for Tokenization 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},