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