epigen/cellwhisperer

1. Overall shape

CellWhisperer is not a single app; it is a repo that combines model code, multiple Snakemake pipelines, dataset curation, evaluation/figure generation, and deployment assets. The main product is a multimodal system for scRNA-seq exploration: transcriptome embeddings plus a chat interface over CELLxGENE.

2. Top-level execution boundaries

The repository separates three major execution modes:

- offline data/model pipelines driven by Snakemake under src/

- Python package runtime under src/cellwhisperer/

- deployment/hosting under hosting/, Dockerfile, and modules/cellxgene

This matters for architecture because most outputs are generated artifacts in resources/ and results/, not source-controlled code.

3. Core embedding package

src/cellwhisperer/ is the main Python library. It contains:

- model definitions for the joint embedding system

- dataset wrappers for training and inference

- loss functions and contrastive objectives

- training entrypoints via PyTorch Lightning

- utilities for inference, model I/O, and processing

This is the central runtime layer used both for training and for generating embeddings for new datasets.

4. Embedding model variants

The package supports multiple transcriptome backbones and variants, including geneformer_model.py, scgpt_model.py, and uce_model.py, plus frozen_model.py and joint embedding orchestration. These are boundary points where the repo adapts external foundation models into a common CellWhisperer interface.

5. Training data engineering

The training corpora are built from several dataset-specific pipelines under src/datasets/ such as ARCHS4/GEO, CELLxGENE Census, Tabula Sapiens, ImmGen, pancreas, bowel disease, and others. Each dataset family has its own Snakemake rules and scripts for downloading, filtering, harmonizing metadata, and producing H5AD-based training inputs.

6. Natural-language annotation pipeline

src/pre_training_processing/ transforms structured single-cell metadata into captions and annotations used for multimodal training. It includes scripts for annotation processing and dataset integration, plus prompts and sample request/response pairs. The README indicates this stage may use GPT-4 when available, or a local Qwen model fallback, so annotation generation is an explicit NLP subpipeline rather than part of the model itself.

7. Joint transcriptome–text training

The contrastive/joint embedding training workflow lives in the cellwhisperer package and is configured through YAML files such as src/cellwhisperer_clip_v1.yaml and the v2 variants. Training is run with the LightningCLI-style cellwhisperer fit ... workflow. Losses and discriminators under src/cellwhisperer/jointemb/loss/ implement the alignment objective between transcriptome and text embeddings.

8. Chat-model pipeline

src/llava/ is a separate training pipeline for the CellWhisperer chat model, built on a modified LLaVA stack. It contains dataset-generation scripts, prompt templates, request/response examples, and Snakemake rules for training. This is architecturally distinct from the embedding model: one subsystem produces vector embeddings, the other powers conversational generation over cell data.

9. User dataset preprocessing for CELLxGENE

src/cellxgene_preprocessing/ is the pipeline that turns a user’s H5AD dataset into a CELLxGENE-ready dataset with CellWhisperer features. It reads raw counts, generates embeddings, annotates clusters, and outputs a processed H5AD used by the browser. This is the main boundary for “bring your own dataset” support.

10. Analysis and paper reproduction

src/figures/ and src/ablation_study/ are downstream analytical pipelines for reproducing manuscript figures, ablations, and validation experiments. They consume outputs from the training and preprocessing stages, often along with precomputed artifacts, and produce the paper’s plots and tables. These directories are analysis consumers rather than production runtime components.

11. Serving and deployment

hosting/ contains the public website and deployment setup, including static pages, nginx configuration, docker compose files, and a dataset-processing Docker setup. src/model_server/ is a lightweight server entrypoint for model-backed functionality. The deployment path also integrates the modified CELLxGENE submodule so the browser can call CellWhisperer-backed services.

12. External dependencies and submodules

modules/ holds git submodules for modified upstream projects: CELLxGENE, LLaVA, Geneformer, and UCE. These are not merely vendored copies; they form explicit architectural seams where the repo extends external ML and browser infrastructure. The repo also depends on external services for hosted AI inference by default, plus optional OpenAI, Hugging Face, and project-hosted dataset/model downloads.

13. Shared infrastructure

src/shared/ centralizes rules and utilities reused across pipelines: model downloads, gene sets, GSVA, sample weights, annotation representations, and transcriptome processing. This layer reduces duplication across dataset prep, training, and evaluation, and is the main cross-cutting dependency among the Snakemake workflows.