CellWorld

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

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.