stargo-cc

STAR-GO checkpoint published for easier discoverability. This repository stores the original Lightning .ckpt and the original TOML config so you can reconstruct the model as trained.

Files

• model.ckpt: PyTorch Lightning checkpoint for TrainingModel
• config.toml: training/model config (same schema as this repo's configs/*.toml)

Provenance

• W&B artifact: contempro-cc-2020-ordered-encdec-medium:best

Usage

This repository contains a Lightning checkpoint and the original TOML config. Load it like this:

import torch
from huggingface_hub import hf_hub_download

from config import from_toml
from model import TrainingModel, get_model_cls

repo_id = "mmtf/stargo-cc"
ckpt_path = hf_hub_download(repo_id, "model.ckpt")
cfg_path = hf_hub_download(repo_id, "config.toml")

cfg = from_toml(cfg_path)

module = TrainingModel.load_from_checkpoint(
    ckpt_path,
    model=get_model_cls(cfg.model.name)(cfg.model),
    training_config=cfg.train,
)
module = module.to("cuda" if torch.cuda.is_available() else "cpu")
module.eval()