File size: 15,253 Bytes
5e449c6
 
 
 
 
 
 
 
 
 
 
 
e1f0461
eaa8608
49ee329
5e449c6
e1f0461
 
5e449c6
e8c145d
e1f0461
 
5e449c6
e1f0461
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eaa8608
e1f0461
 
 
ba359cc
 
 
e1f0461
 
 
07bb78a
e1f0461
 
 
3830179
 
 
 
 
 
 
 
 
 
 
e8c145d
e1f0461
e8c145d
 
 
e1f0461
 
 
eaa8608
 
e1f0461
 
 
eaa8608
e1f0461
 
eaa8608
5ed2951
 
 
3830179
 
 
e1f0461
 
 
 
 
 
 
 
 
999e5eb
5e449c6
3830179
 
 
 
 
e1f0461
3ecef9f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49ee329
e1f0461
49ee329
e1f0461
 
 
eaa8608
e1f0461
faf1f6d
 
 
 
e1f0461
 
 
 
 
 
 
 
 
 
3830179
 
 
 
 
 
e1f0461
 
 
faf1f6d
 
 
e8c145d
eaa8608
e1f0461
 
 
 
 
 
faf1f6d
 
 
e1f0461
 
 
 
 
 
 
49ee329
e1f0461
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5e449c6
 
e1f0461
 
 
 
 
 
5e449c6
 
3830179
 
 
 
 
 
 
5e449c6
3830179
 
 
5e449c6
3830179
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5e449c6
3830179
5e449c6
 
3830179
 
5e449c6
 
eaa8608
5e449c6
 
 
3830179
 
eaa8608
 
 
3830179
 
 
 
 
5e449c6
 
eaa8608
 
5e449c6
5ed2951
5e449c6
 
 
3830179
eaa8608
 
5e449c6
 
 
 
 
 
 
 
 
e1f0461
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
---
license: other
tags:
  - single-cell
  - transcriptomics
  - foundation-model
  - circulating-tumor-cells
library_name: pytorch
---

# BioMod

BioMod is a single-cell transcriptomics foundation model built around gene modules
and an attention-based encoder. This repository provides pretrained model assets,
pretraining and fine-tuning code, and downstream task tutorials.

BioMod is released by **CTC-Atlas**, the joint laboratory of **Raymind Biotechnology
Co., Ltd** and the **Shenzhen Research Institute of Big Data (SRIBD)**.

[Installation](#installation) · [BioMod model weights](#biomod-model-weights) ·
[Tasks and examples](#tasks-and-examples) · [Usage notes](#usage-notes) ·
[Repository layout](#repository-layout)

## Installation

Download the project from [Hugging Face](https://huggingface.co/ctcatlas/BioMod).
Run the following commands from the parent of your intended project directory:

```bash
python -m pip install --upgrade huggingface_hub
hf auth login
hf download ctcatlas/BioMod --local-dir ./BioMod
cd BioMod
```

The download includes source code, model dependencies, the pretrained checkpoint,
and datasets. To omit datasets, add `--exclude "data/**"` to the download command.
If you already have the project, start from its root directory.

Create and activate the supplied Conda environment:

```bash
conda env create -f BioMod/environment.yml
conda activate biomod
```

The environment specification uses Python 3.9, PyTorch 2.6, CUDA 12.4, and PyTorch
Geometric 2.6. Conda creates the Python environment, then installs the pinned
GPU packages from the PyTorch/PyG wheel indexes through `requirements.txt`.
See [environment.yml](BioMod/environment.yml) and
[requirements.txt](BioMod/requirements.txt) for the complete dependency lists.
Match the GPU packages to your system when adapting the environment.

Pretraining uses CUDA/NCCL. The zero-shot
runner selects CUDA when available and otherwise uses CPU. Native Windows
execution has not been validated for the supplied environment.

To open a tutorial, choose its task directory. For cell annotation, run the
following from the outer repository root:

```bash
cd BioMod/cell_annotation
jupyter lab tutorial.ipynb
```

Use `drug_response`, `drug_combination`, `gene_perturbation`, or `zero_shot_batch`
in place of `cell_annotation` to open another task.

## BioMod model weights

This release includes the **BioMod model weights** and their matching
configuration. Use these files together; the configuration records the model
architecture and gene-module definitions.

| Asset | File |
| --- | --- |
| BioMod model weights | [BioMod.pth](weight/BioMod.pth) |
| Matching model configuration | [BioMod.pkl](weight/BioMod.pkl) |
| Model gene order | [selected_genes_27k.txt](model_dependencies/selected_genes_27k.txt) |
| Gene identifier mapping | [gene2ens.csv](model_dependencies/gene2ens.csv) |
| Gene2Vec embeddings | [selected_gene2vec_27k.npy](model_dependencies/selected_gene2vec_27k.npy) |
| Shared module definitions | [merged_biomods.json](model_dependencies/biomod_data_all/merged_biomods.json) |

The model components are bundled under
[`model_dependencies/pytorch_cont/`](model_dependencies/pytorch_cont/).
Legacy serialized names are handled by
[`biomod_compat.py`](model_dependencies/biomod_compat.py) when loading existing
checkpoints; the original checkpoint files remain unchanged.
Optional [Hallmark and Reactome source definitions](model_dependencies/biomod_data_sources/)
are kept in `biomod_data_sources/`. They are separate from the default checkpoint
runtime definitions in `biomod_data_all/`.
Task-specific checkpoint arguments are listed under
[Usage notes](#checkpoint-arguments).

## Tasks and examples

The following tasks have implementations in this repository. The links identify
available source code and examples; these workflows have not been validated end
to end with the bundled checkpoint.

| Task | Implementation | Notebook |
| --- | --- | --- |
| Cell-type annotation | [Fine-tuning script](BioMod/cell_annotation/train.py) | [Notebook](BioMod/cell_annotation/tutorial.ipynb) |
| Zero-shot cell embeddings and batch-effect evaluation | [Embedding and evaluation script](BioMod/zero_shot_batch/evaluate.py) | [Notebook](BioMod/zero_shot_batch/tutorial.ipynb) |
| Single-drug response regression (IC50) | [Training script](BioMod/drug_response/train.py) | [Notebook](BioMod/drug_response/tutorial.ipynb) |
| Drug-pair synergy classification | [Training script](BioMod/drug_combination/train.py) | [Notebook](BioMod/drug_combination/tutorial.ipynb) |
| Gene perturbation response prediction | [Training script](BioMod/gene_perturbation/train.py) | [Notebook](BioMod/gene_perturbation/tutorial.ipynb) |

### Commands by dataset

Each row below links to a separate, complete launch command for that task and
dataset or condition. Commands specify the input paths and output directory;
you do not need to adapt a command from another dataset. Run them in Bash from
the outer repository root after activating the environment. Read any
dataset-specific input caveats before launching.

| Task | Dataset or condition | Launch command |
| --- | --- | --- |
| Cell annotation | Kim Lung | [Complete command](BioMod/cell_annotation/README.md#1-kim-lung) |
| Cell annotation | Immune | [Complete command](BioMod/cell_annotation/README.md#2-immune) |
| Cell annotation | Liver | [Complete command](BioMod/cell_annotation/README.md#3-liver) |
| Cell annotation | MS (paired train/test files) | [Complete command](BioMod/cell_annotation/README.md#4-ms) |
| Cell annotation | hPancreas (hp14 paired files) | [Complete command](BioMod/cell_annotation/README.md#5-hpancreas) |
| Cell annotation | Zheng68K | [Complete command](BioMod/cell_annotation/README.md#6-zheng68k) |
| Cell annotation | Myeloid (reference/query files) | [Complete command](BioMod/cell_annotation/README.md#7-myeloid) |
| Zero-shot evaluation | hPNC | [Complete command](BioMod/zero_shot_batch/README.md#1-hpnc) |
| Zero-shot evaluation | BMMC | [Complete command](BioMod/zero_shot_batch/README.md#2-bmmc) |
| Zero-shot evaluation | Kim Lung: malignant cells retained | [Complete command](BioMod/zero_shot_batch/README.md#3-kim-lung-malignant-cells-retained) |
| Zero-shot evaluation | Kim Lung: malignant cells removed | [Complete command](BioMod/zero_shot_batch/README.md#4-kim-lung-malignant-cells-removed) |
| Single-drug response | GDSC-CCLE joint inputs | [Complete command](BioMod/drug_response/README.md#run-the-bundled-gdsc-ccle-collection) |
| Drug-combination synergy | Bundled training and validation collection | [Complete command](BioMod/drug_combination/README.md#run-the-bundled-training-and-validation-collection) |
| Gene perturbation | Norman | [Complete command](BioMod/gene_perturbation/README.md#2-run-the-bundled-norman-simulation-split) |

Paired training/test files and complementary drug/cell inputs belong to one
run configuration. Historical reference copies do not add datasets. The
[optional drug-combination cross-validation command](BioMod/drug_combination/README.md#optional-five-fold-cross-validation-on-the-same-collection)
uses the same collection under a different validation protocol.

### Pretraining

The repository also supports pretraining and continued pretraining.

| Workflow | Implementation | Documentation |
| --- | --- | --- |
| Pretraining and continued pretraining | [Pretraining script](BioMod/biomod.py) | [Source guide](BioMod/README.md), [launcher reference](BioMod/command.txt) |

Bundled input defaults resolve from each script's location. Command examples
state their working directory; explicit relative command-line paths are resolved
from that directory. Pretraining corpora, legacy reference checkpoints, and
generated analysis results must be supplied separately where indicated.

## Usage notes

### Project paths and notebooks

The project root contains `BioMod/`, `model_dependencies/`, `weight/`, and
`data/`. The inner `BioMod/` directory contains the source code. Task commands
may assume different working directories; adjust relative paths when following
a guide.

Each task keeps its notebook in `BioMod/<task>/tutorial.ipynb`. Start Jupyter
from the selected task directory, as shown in [Installation](#installation).
Each notebook sets `PROJECT_ROOT = Path('../..').resolve()` and derives absolute
runtime paths from the outer repository root. If the notebook kernel uses a
different working directory, set `PROJECT_ROOT` to the outer repository root
explicitly before running the path checks. Review the configuration cells before running:

| Setting | Location |
| --- | --- |
| `PROJECT_ROOT` | Outer project root containing `BioMod/`, `model_dependencies/`, `weight/`, and `data/` |
| `SOURCE_DIR` | `PROJECT_ROOT / 'BioMod'` |
| Asset or dependency directory | `PROJECT_ROOT / 'model_dependencies'` |
| Weight and configuration paths | The BioMod model weights and matching configuration listed above |
| `MODULE_DATA_DIR` | `model_dependencies/biomod_data_all/` under the project root |
| Data and output paths | Your task inputs and a writable output directory |

Checkpoint variable names differ by notebook: `CKPT_PATH` / `CONFIG_PKL`
for cell annotation, drug response, and zero-shot integration;
`SCFM_CKPT` / `SCFM_HYPER` for gene perturbation; and
`CHECKPOINT_PATH` / `HYPER_PARAMS_PATH` for drug combination.
Training defaults to disabled (`RUN_TRAINING = False`). The zero-shot notebook
uses `RUN_EVALUATION = False` and enables UMAP export separately with `MAKE_UMAP`.
Output directories are created when the corresponding operation is enabled.

### Checkpoint arguments

| Task | Weights argument | Configuration argument |
| --- | --- | --- |
| Drug response / drug combination | `--checkpoint_path` | `--hyper_params_path` |
| Zero-shot / gene perturbation | `--scfm_ckpt_path` | `--scfm_hyper_params_path` |
| Cell annotation | `--biomod_checkpoint` | `--config_pkl` |
| Continued pretraining | `--model_path` | `--model_config_path` or paired-file discovery |

### Data and evaluation

Datasets are stored under `data/`. Input formats, metadata fields, preprocessing,
and train/validation splits depend on the task; follow the linked implementation
and tutorial.

For zero-shot embeddings, provide an h5ad expression matrix and specify the gene,
batch, and cell-type metadata fields. The runner aligns genes to the model gene
list and fills missing genes with zeros. It does not normalize or log-transform
the input matrix, so prepare expression values according to the intended protocol.
Resolve duplicate gene identifiers and inspect gene overlap before evaluating
embeddings.

The drug-response and drug-combination workflows use validation metrics for
checkpoint selection. Their validation results should not be presented as
performance on an independent final test set.

## Repository layout

The tree below includes source files, model dependencies, and the pretrained
checkpoint. The top-level `data/` and `.cache/` directories are omitted for
readability. Folder descriptions explain where to find each component.

<details>
<summary>Expand the full repository structure</summary>

```text
BioMod/                                                        # Project root: source code, model dependencies and weights
├── BioMod/                                                    # Task code: training, evaluation and tutorials
│   ├── cell_annotation/                                       # Cell-type annotation and model fine-tuning
│   │   ├── README.md
│   │   ├── train.py
│   │   └── tutorial.ipynb
│   ├── drug_combination/                                      # Drug-combination synergy prediction
│   │   ├── README.md
│   │   ├── train.py
│   │   └── tutorial.ipynb
│   ├── drug_response/                                         # Single-drug response prediction
│   │   ├── README.md
│   │   ├── train.py
│   │   └── tutorial.ipynb
│   ├── gene_perturbation/                                     # Gene perturbation response prediction
│   │   ├── components/                                        # Task models, datasets, training and inference helpers
│   │   │   ├── __init__.py
│   │   │   ├── data_utils.py
│   │   │   ├── dataset.py
│   │   │   ├── inference.py
│   │   │   ├── model.py
│   │   │   ├── trainer.py
│   │   │   ├── training_utils.py
│   │   │   └── utils.py
│   │   ├── README.md
│   │   ├── train.py
│   │   └── tutorial.ipynb
│   ├── zero_shot_batch/                                       # Zero-shot integration and evaluation
│   │   ├── tools/                                             # Optional visualization and evaluation tools
│   │   │   ├── sequencing_depth/                              # Low-depth expression generation and evaluation
│   │   │   │   ├── make_low_depth_h5ad.py
│   │   │   │   ├── README.md
│   │   │   │   └── run_zero_shot_low_depth.py
│   │   │   └── visualize_results.py
│   │   ├── evaluate.py
│   │   ├── README.md
│   │   └── tutorial.ipynb
│   ├── .gitattributes
│   ├── .gitignore
│   ├── biomod.py
│   ├── biomod_dataset.py
│   ├── command.txt
│   ├── environment.yml
│   ├── models_biomod.py
│   ├── README.md
│   ├── requirements.txt
│   └── training_utils.py
├── model_dependencies/                                        # Shared gene mappings, embeddings and model components
│   ├── biomod_data_all/                                       # Default merged module definitions and construction tools
│   │   ├── create_uncovered_biomods.py
│   │   ├── merge_biomods.py
│   │   └── merged_biomods.json
│   ├── biomod_data_sources/                                   # Optional source gene sets and their loading helper
│   │   ├── biomod1_hallmarks.json
│   │   ├── biomod2_reactome_2022_filtered.json
│   │   └── load_biomod.py
│   ├── pytorch_cont/                                          # Attention and model components
│   │   ├── __init__.py
│   │   ├── ding_models.py
│   │   ├── pytorch.py
│   │   ├── pytorch_old.py
│   │   └── reversible.py
│   ├── biomod_compat.py                                      # Legacy checkpoint and configuration compatibility
│   ├── gene2ens.csv
│   ├── selected_gene2vec_27k.npy
│   └── selected_genes_27k.txt
├── weight/                                                    # Pretrained weights and matching model configuration
│   ├── BioMod.pkl
│   └── BioMod.pth
├── .gitattributes
├── .gitignore
└── README.md
```

</details>

## Partners

- [Raymind Biotechnology Co., Ltd](https://en.raymind.com/)
- [Shenzhen Research Institute of Big Data (SRIBD)](https://www.sribd.cn/en)