ProteinGym Dataset

The ProteinGym dataset implementation. ProteinGym is a large-scale benchmark for protein fitness prediction, containing multiple protein families with experimental fitness measurements.

class alf_tools.datasets.proteingym.ProteinGym(config)

Bases: BaseDataset

ProteinGym dataset class.

load_dataset()

Load ProteinGym dataset, downloading from the public ProteinGym_v1 HF repository if a local cache is not present.

Data is sourced from OATML-Markslab/ProteinGym_v1 — no HuggingFace token is required. CV fold columns are computed deterministically and saved alongside the raw data so subsequent loads are instant.

Note

The local cache at data/ProteinGym/<dms_name>/data.csv is not automatically invalidated when the upstream dataset is updated. Set force_download=True in the config to delete the cache and re-download, or remove the file manually.

Return type:

LabelledCandidates

Returns:

Labeled candidates with ProteinGym data.

class alf_tools.datasets.proteingym.ProteinGymConfig(**data)

Bases: BaseDatasetConfig

Configuration for ProteinGym dataset.

dms_name

Name of the DMS assay (e.g., “IF1_ECOLI_Kelsic_2016”).

dms_type

Type of DMS data (“singles” or “multiples”).

cross_validation

Whether to use cross-validation splits.

cross_validation_type

Type of CV split (“random”, “modulo”, or “contiguous”). Only “random” folds are available for dms_type=”multiples”.

cross_validation_fold

Which CV fold to use (0-4).

force_download

If True, delete the local cache before loading so the dataset is re-downloaded from upstream. Use this to pick up corrected versions of an assay published by ProteinGym_v1.

cross_validation: bool

cross_validation_fold: Literal[0, 1, 2, 3, 4] | None

cross_validation_type: Literal['random', 'modulo', 'contiguous'] | None

dms_name: str

dms_type: Literal['singles', 'multiples']

force_download: bool

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

problem_type: ProblemType