library_name: pytorch
license: mit
pipeline_tag: other
tags:
- spatial-transcriptomics
- self-supervised-learning
- jepa
CellWorld pretrained models
CellWorld is introduced in the paper CellWorld: From Gene-Level Reconstruction to Latent Cell Prediction in Spatial Transcriptomics Foundation 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.
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
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.