KaleProtein

Reusable components and consistent interfaces for protein AI.

Installation · Quick Start · Models and Datasets · Model and Data Cards

KaleProtein is a modular toolkit for developing and using protein AI models across multiple modalities. It brings reusable neural network components, dataset adapters, and model cards together through consistent Auto APIs.

• For protein AI model developers: Compose existing protein and multimodal components, develop and test models against explicit interfaces, and share implementations through a common model-card layout.
• For protein scientists: Load datasets and contributed models through Auto APIs, run predictions or generate protein sequences, and evaluate models with reusable metrics.

Data preparation, embedding, prediction or generation, evaluation, and interpretation remain explicit steps. Named dictionaries connect them, so researchers can inspect intermediate results and replace individual components.

The data loader owns data preparation and batching. The model owns its embedders, predictor or generator, and model weights. Auto APIs select and construct registered implementations. See the architecture guide for component boundaries and input/output contracts.

Requires Python 3.10 or later.

Choose the installation that fits your work:

Replace example-name with an example name from the model index. Use the same name when installing its dependencies and downloading its code. The dev extra does not install example-specific dependencies.

The wheel installs the kaleprotein library. Extras add dependencies; they do not install the repository's examples/, tests/, or docs/.

Install the selected example's dependencies, then download its code. For example:

python -m pip install "kaleprotein[drugban_dti]"
python -m kaleprotein download-example drugban_dti

Examples are saved under ./examples/<name>/. Run the snippets below from the directory containing examples/; no library checkout or editable install is needed. The downloader fetches code, configuration, maps, and license notices, without importing the downloaded code. It leaves existing example directories untouched.

By default, examples come from the installed library's v<version> release tag. Use --ref <tag-or-commit> to select a revision, or --ref main when working with development code. --output <directory> changes the parent download directory. The selected commit is recorded in .kaleprotein-example.json inside each downloaded example.

Datasets and large checkpoints are obtained separately. The model and dataset index links to each workflow's requirements.

This DrugBAN example evaluates one batch of drug-target pairs and exposes interaction attention. Prepare the DrugBAN-formatted BindingDB CSV files and a trained DrugBAN checkpoint first, following the DrugBAN guide. Replace the two local paths below with your dataset root and checkpoint.

import torch

from kaleprotein.auto import (
    AutoProteinConfig,
    AutoProteinDataLoader,
    AutoProteinInterpreter,
    AutoProteinModel,
)
from examples.drugban_dti import register_model_card

register_model_card()
config = AutoProteinConfig.from_pretrained("DTI/DrugBAN")

# 1. Load, preprocess, and batch the dataset.
loader = AutoProteinDataLoader(
    "BindingDB/DTI",
    config=config,
    root="path/to/DrugBAN/datasets",
    split="random",
    subset="test",
    batch_size=64,
)
inputs = next(iter(loader))

# 2. Build the model and load its checkpoint.
model = AutoProteinModel.from_config(config, checkpoint="path/to/drugban.pt")
model.eval()

with torch.inference_mode():
    # 3. Embed the protein and molecule.
    embeddings = model.embed(**inputs)

    # 4. Predict interactions.
    prediction = model.predict(**embeddings)

    # 5. Evaluate the predictions.
    metrics = model.evaluate(**prediction)

    # 6. Optionally interpret interaction attention.
    interpreter = AutoProteinInterpreter.from_config(config)
    interpretation = interpreter.explain(**prediction)

print(metrics)

Stage outputs are dictionaries with named fields. For example, DrugBAN embeddings include protein_embedding, molecule_embedding, and their masks. Labels and sample metadata travel alongside them.

This snippet reports one-batch metrics. Use the full evaluation workflow to evaluate a complete split. Model comparisons should use the same held-out records and metric definitions, with each model's required preprocessing.

Generative models follow the same pattern. After constructing the MapDiff loader and model, encode the structure and generate sequences:

with torch.inference_mode():
    embeddings = model.embed(**inputs)
    generation = model.generate(**embeddings, steps=100, method="ddim")
    metrics = model.evaluate(**generation)

Here, model and inputs refer to MapDiff, not the DrugBAN objects above. See the complete MapDiff pipeline for setup, pretrained weights, and structure-input requirements.

Each module has a defined role and provides reusable pipeline components. The embed and predict modules live under kaleprotein.model:

AutoProteinDataLoader combines loaddata and prepdata: records are loaded, preprocessed, and then collated into batches. AutoProteinModel combines the embedders and predictor or generator and handles checkpoint loading. The model consumes prepared batches; it does not own the dataset, preprocessor, or collator.

Evaluation and interpretation are independent consumers of model outputs. Use model.evaluate(**outputs) or a configured AutoProteinEvaluator for metrics, and AutoProteinInterpreter for optional explanations. Interpretation does not require running evaluation first.

Auto APIs select and construct registered components from configuration. The dictionary keys passed through **inputs, **embeddings, and **outputs define the interfaces between stages, allowing developers to choose their own fields and replace compatible components. Shared utilities provide supporting operations such as file parsing and checkpoint handling.

See the extension guide for implementing and registering custom pipeline components and model cards.

Complete model implementations live in the repository examples. Both are self-contained PyTorch refactors and do not import an upstream checkout at runtime.

AutoProteinModel(..., pretrain=True) checks the card's local weight path, downloads from its configured URL when needed, and verifies a checksum when provided. If no local checkpoint or valid URL exists, it raises an error explaining that training or a supplied checkpoint is required. Use checkpoint="path/to/model.pt" to load a specific local checkpoint.

Dataset IDs follow Dataset/Task. Adapters load local data and normalize records for reuse across compatible models.

Use AutoProteinData for dataset records alone, or AutoProteinDataLoader to combine dataset loading with a configured preprocessor and collator:

from kaleprotein.auto import AutoProteinData

data = AutoProteinData(
    "BindingDB/DTI",
    root="path/to/DrugBAN/datasets",
    split="random",
    subset="test",
)

DTI adapters target the processed benchmark CSVs, not arbitrary raw exports from the original databases. For MapDiff, processed CATH graphs preserve training features that the raw PDB/mmCIF path cannot fully reconstruct; see the preparation notes before comparing results.

Cards make implementations discoverable and document the inputs, configuration, and assets needed to use them.

Model cards: The current layout includes config.yaml, a configuration class, the complete model implementation, and local asset folders. The config declares Auto class mappings, preprocessing streams, reusable components, and pretrained-weight metadata. See the model-card format and implementation guide and external registration instructions.

Data cards: A standalone data-card file format is not implemented yet. Datasets currently use registered Python adapters with Dataset/Task IDs. They normalize task fields and preserve provenance independently of any model. The dataset extension guide documents this current interface.

Models can be shared through repository contributions or external card directories that users register locally. A hosted upload service is not part of the current library.

Contributions can add reusable components, dataset adapters, complete model cards, tests, or documentation. Use the extension guide for implementation conventions and open a pull request against this repository.

For local development:

git clone https://github.com/pykale/protein.git
cd protein
python -m pip install -e ".[dev]"
python -m pytest -q

Install the extra for any example you are developing separately, using its name from the model index.

Tests use temporary data, mocked download clients, and small model inputs; they do not download real pretrained checkpoints.

Report bugs and discuss proposed additions through GitHub Issues.

When using a contributed model or dataset in research, cite its original work as described in the corresponding example and source documentation. Upstream attribution and license details are collected in THIRD_PARTY_NOTICES.md.

KaleProtein is released under the MIT License.