| --- |
| 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. |
|
|
