weililab/perturbation-embeddings

Gene embeddings for perturbation modeling

This repository distributes three existing gene embedding tables in a common NumPy/JSON format for use with pertTF and other perturbation models. These are gene-level features, not measured perturbation responses or pertTF-generated predictions. The original embedding models were not trained by this release.

Tables

Each directory contains:

• —embeddings.npy: a dense matrix with shape (n_genes, embedding_dim).
• —genes.json: an ordered JSON list; genes[i] labels row i in the matrix.
• —metadata.json: source identifiers, processing details, file hashes, and conversion verification results.

ESM2 and GenePT are restricted to gene symbols in the first column of their respective human_prot_names.csv files. Only listed genes present in each embedding table are retained, preserving original gene identifiers and row order. Gene-list hashes, matched counts, and unmatched counts are recorded in each table's metadata. GEARS is unchanged. No normalization, identifier conversion, concatenation, or synthetic control (WT) row was added.

This corrects the initial release, which included the complete unfiltered GenePT mapping. The earlier revision 5db008c45c77caaf5ce9e7c9783c89ce1d8a2f25 remains available for reproducing that release.

The ESM2 and GEARS inputs were float32. The GenePT input was float64, but its values were exactly representable in float32: the maximum absolute conversion error was zero for all three tables. Every exported row was checked against its corresponding source vector after reloading the output files, and the saved gene ordering was verified.

Loading

These are ordinary files hosted in an HF dataset repository; the datasets package is not required.

import json
import numpy as np
from huggingface_hub import HfApi, hf_hub_download

repo_id = "weililab/perturbation-embeddings"
embedding = "esm2"  # "genept" or "gears"
# Resolve once so the matrix and gene list come from the same snapshot.
# Save this commit SHA with your experiment, or supply a previously saved SHA.
revision = HfApi().dataset_info(repo_id).sha

matrix_path = hf_hub_download(
    repo_id, f"{embedding}/embeddings.npy",
    repo_type="dataset", revision=revision,
)
genes_path = hf_hub_download(
    repo_id, f"{embedding}/genes.json",
    repo_type="dataset", revision=revision,
)
matrix = np.load(matrix_path, mmap_mode="r", allow_pickle=False)
with open(genes_path) as handle:
    genes = json.load(handle)
assert matrix.shape[0] == len(genes)

Gene coverage differs between tables. Align by gene identifier rather than row position when comparing or concatenating representations. A trained model's input representation cannot generally be replaced with another embedding choice without retraining.

Sources, attribution, and licenses

Licensing is recorded per source; no single license is asserted for all three tables.

ESM2 / UCE

• —Source record: Yusuf Roohani (2023), Universal Cell Embedding Model Files, Figshare, version 5: https://doi.org/10.6084/m9.figshare.24320806.v5.
• —The record distributes protein_embeddings.tar.gz and specifies CC BY 4.0.
• —The local input was Homo_sapiens.GRCh38.gene_symbol_to_embedding_ESM2.pt, associated with that UCE download. This release did not byte-compare the individual file against the original archive.
• —Credit the UCE authors and the ESM2 authors when using these embeddings. ESM2 reference: Lin et al. (2023), Evolutionary-scale prediction of atomic-level protein structure with a language model, https://doi.org/10.1126/science.ade2574.

GenePT

• —Source record: Yiqun Chen (2024), Gene embeddings used in GenePT, Zenodo: https://doi.org/10.5281/zenodo.10833191.
• —The record specifies CC BY 4.0 and identifies the model-3 embedding as text-embedding-3-large.
• —The input filename, GenePT_gene_protein_embedding_model_3_text.pickle, matches the file described in GenePT_emebdding_v2.zip. This release did not byte-compare the local file against the original archive.
• —Please cite Chen and Zou, GenePT: A Simple But Effective Foundation Model for Genes and Cells Built From ChatGPT: https://doi.org/10.1101/2023.10.16.562533.
• —Project: https://github.com/yiqunchen/GenePT.

GEARS

• —Extracted pert_emb.weight from the legacy checkpoint gears_gwps-Mar25-03-38/best_model/model.pt, using the ordered pert_gene_list from its accompanying pert_gene_list.pkl.
• —Training provenance is unknown. The training dataset, split, and exact training configuration were not reconstructed for this release.
• —The checkpoint's license is unknown; the CC BY 4.0 licenses of the other sources are not attributed to this checkpoint.
• —GEARS project: https://github.com/snap-stanford/GEARS.

The changes made by this release are human-gene filtering, format conversion, and float32 storage, as described above. Original-file SHA-256 hashes are recorded in each table's metadata to identify the exact local inputs used.