File size: 2,074 Bytes
dd89720 0dc45de dd89720 | 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 | ---
license: mit
library_name: pytorch
tags:
- spatial-transcriptomics
- self-supervised-learning
- jepa
---
# CellWorld pretrained models
This repository contains the four canonical CellWorld pretrained
models used for model-scale experiments. CellWorld is a joint-embedding
predictive architecture for spatial transcriptomics.
## Checkpoints
| Scale | Encoder | Heads | Predictor | Training steps |
|---|---:|---:|---:|---:|
| `small` | 192 x 6 | 6 | 96 x 3 | 10,000 |
| `base` | 384 x 8 | 6 | 192 x 3 | 10,000 |
| `large` | 512 x 12 | 8 | 256 x 3 | 10,000 |
| `huge` | 768 x 16 | 12 | 384 x 3 | 10,000 |
`base` is the recommended default. Each folder contains:
- `model.pt`: complete CellWorld model state and the minimal architecture
configuration required by the released code.
- `config.yaml`: the corresponding path-free pretraining recipe.
The release checkpoints intentionally omit optimizer state, data locations,
output locations, resume paths, and experiment-service metadata. They retain
the complete pretrained model, including the context encoder, target encoder,
cell tokenizer, metadata embeddings, and predictor.
## Download
Authentication is required because this model repository is private.
```python
from huggingface_hub import hf_hub_download
checkpoint = hf_hub_download(
repo_id="Haiping-UoM/CellWorld",
filename="base/model.pt",
token=True,
)
print(checkpoint)
```
The downloaded checkpoint can be passed directly to the `--pretrained-ckpt`
argument of the CellWorld probe and finetuning commands.
## Checkpoint schema
```text
format_version Release checkpoint format version
scale small, base, large, or huge
model Complete CellWorld state_dict
config Model and vocabulary configuration used for loading
epoch Completed training epoch
step Completed optimizer step
```
Checksums, byte sizes, and tensor-element counts are recorded in
`manifest.json`.
## Code and license
Source code: https://github.com/UoM-HealthAI/CellWorld
Released under the MIT License.
|