kernel-elastic-autoencoder 1.0.0__tar.gz

kernel_elastic_autoencoder-1.0.0/PKG-INFO

@@ -0,0 +1,44 @@
+Metadata-Version: 2.4
+Name: kernel_elastic_autoencoder
+Version: 1.0.0
+Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
+License: MIT
+Author: Felix Rotter-McCartney
+Author-email: felix.rotter@mail.utoronto.ca
+Requires-Python: >=3.12,<3.15
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: huggingface-hub (>=1.22.0,<2.0.0)
+Requires-Dist: pandas (>=3.0.3,<4.0.0)
+Requires-Dist: pydantic (>=2.13.4,<3.0.0)
+Requires-Dist: safetensors (>=0.8.0,<0.9.0)
+Description-Content-Type: text/markdown
+
+## `kernel_elastic_autoencoder`
+
+`kernel_elastic_autoencoder` is a library implementing the architecture and techniques described
+in [this publication](https://doi.org/10.1093/pnasnexus/pgae168) from Li et al. I am not affiliated with the authors of
+the original paper, and this implementation is provided as-is, with no guarantee of completeness.
+
+### Installation
+
+`kernel_elastic_autoencoder` can be installed with pip, and regular builds are provided on PyPI:
+
+    pip install kernel_elastic_autoencoder
+
+> Please note that `torch` is not included as a dependency due to its many hardware-accelerator-dependent versions, so
+> take care to install the appropriate version manually.
+
+Distribution builds are also provided here on GitHub Releases. New builds are triggered by the CD Action, so they will 
+be made available as soon as a new PR is merged to `main`.
+
+Alternatively, for development purposes, `kernel_elastic_autoencoder` may be installed from source provided here. Builds
+and deps are managed with `poetry`.
+
+### Documentation
+
+API documentation is generated with `pdoc` and covers the `__all__`-exported interfaces. It is available on 
+[GitHub Pages](https://cancelradius.github.io/kernel_elastic_autoencoder).

kernel_elastic_autoencoder-1.0.0/README.md

@@ -0,0 +1,25 @@
+## `kernel_elastic_autoencoder`
+
+`kernel_elastic_autoencoder` is a library implementing the architecture and techniques described
+in [this publication](https://doi.org/10.1093/pnasnexus/pgae168) from Li et al. I am not affiliated with the authors of
+the original paper, and this implementation is provided as-is, with no guarantee of completeness.
+
+### Installation
+
+`kernel_elastic_autoencoder` can be installed with pip, and regular builds are provided on PyPI:
+
+    pip install kernel_elastic_autoencoder
+
+> Please note that `torch` is not included as a dependency due to its many hardware-accelerator-dependent versions, so
+> take care to install the appropriate version manually.
+
+Distribution builds are also provided here on GitHub Releases. New builds are triggered by the CD Action, so they will 
+be made available as soon as a new PR is merged to `main`.
+
+Alternatively, for development purposes, `kernel_elastic_autoencoder` may be installed from source provided here. Builds
+and deps are managed with `poetry`.
+
+### Documentation
+
+API documentation is generated with `pdoc` and covers the `__all__`-exported interfaces. It is available on 
+[GitHub Pages](https://cancelradius.github.io/kernel_elastic_autoencoder).

kernel_elastic_autoencoder-1.0.0/pyproject.toml

@@ -0,0 +1,46 @@
+[project]
+name = "kernel_elastic_autoencoder"
+version = "1.0.0"
+description = "Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)"
+authors = [
+    { name = "Felix Rotter-McCartney", email = "felix.rotter@mail.utoronto.ca" }
+]
+readme = "README.md"
+license = { text = "MIT" }
+requires-python = ">=3.12,<3.15"
+dependencies = [
+    "pydantic (>=2.13.4,<3.0.0)",
+    "huggingface-hub (>=1.22.0,<2.0.0)",
+    "pandas (>=3.0.3,<4.0.0)",
+    "safetensors (>=0.8.0,<0.9.0)",
+]
+
+[tool.poetry]
+packages = [{ include = "kernel_elastic_autoencoder", from = "src" }]
+
+[tool.semantic_release]
+version_toml = ["pyproject.toml:project.version"]
+commit_parser = "conventional"
+build_command = "pip install build && python -m build --sdist --wheel ."
+
+[tool.semantic_release.branches.main]
+match = "(main|master)"
+prerelease = false
+
+[tool.semantic_release.remote.token]
+env = "GH_TOKEN"
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[dependency-groups]
+dev = [
+    "pytest (>=9.1.1,<10.0.0)",
+    "ruff (>=0.15.20,<0.17.0)",
+    "pytest-cov (>=7.1.0,<8.0.0)",
+    "pytest-xdist (>=3.8.0,<4.0.0)",
+    "ty (>=0.0.62,<0.0.64)",
+    "python-semantic-release (>=10.6.1,<11.0.0)",
+    "pdoc (>=16.0.0,<17.0.0)",
+]

kernel_elastic_autoencoder-1.0.0/src/kernel_elastic_autoencoder/__init__.py

@@ -0,0 +1,44 @@
+from kernel_elastic_autoencoder.collate import Collated, Collator, DataframeCollator
+from kernel_elastic_autoencoder.config import (
+    ExperimentConfig,
+    ModelCommonConfig,
+    ModelConfig,
+    ModelDecoderConfig,
+    ModelEncoderConfig,
+    ModelInputConfig,
+    TrainingCommonConfig,
+    TrainingConfig,
+    TrainingHyperparameterConfig,
+    TrainingOptimizerConfig,
+)
+from kernel_elastic_autoencoder.losses import Loss
+from kernel_elastic_autoencoder.model import Model
+from kernel_elastic_autoencoder.pipeline import Completion, Pipeline
+from kernel_elastic_autoencoder.sample import Sampler, Top1Sampler
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+from kernel_elastic_autoencoder.training import Trainer, TrainerCallback
+
+__all__ = [
+    "ExperimentConfig",
+    "ModelConfig",
+    "ModelCommonConfig",
+    "ModelInputConfig",
+    "ModelEncoderConfig",
+    "ModelDecoderConfig",
+    "TrainingConfig",
+    "TrainingCommonConfig",
+    "TrainingHyperparameterConfig",
+    "TrainingOptimizerConfig",
+    "Pipeline",
+    "Completion",
+    "Model",
+    "Loss",
+    "Tokenizer",
+    "Collator",
+    "Collated",
+    "DataframeCollator",
+    "Sampler",
+    "Top1Sampler",
+    "Trainer",
+    "TrainerCallback",
+]

kernel_elastic_autoencoder-1.0.0/src/kernel_elastic_autoencoder/collate.py

@@ -0,0 +1,164 @@
+from abc import ABC, abstractmethod
+
+import torch
+from pandas import DataFrame
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from kernel_elastic_autoencoder.config import ModelConfig
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+
+
+class Collated(BaseModel):
+    """Return schema for Collator.__call__. Represents a store of tensors ready for direct input to a Model."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    input_ids: torch.Tensor = Field(
+        ...,
+        description="Tensor of dimension (B, S) containing indices of input tokens, generated by a Tokenizer.",
+    )
+    conditions: torch.Tensor = Field(
+        ...,
+        description="Tensor of dimension (B, C) containing condition values for each sequence.",
+    )
+    token_mask: torch.Tensor = Field(
+        ...,
+        description="Tensor of dimension (B, S) containing a boolean padding mask for each sequence.",
+    )
+    condition_mask: torch.Tensor = Field(
+        ...,
+        description="Tensor of dimension (B, C) containing a boolean condition mask for each sequence",
+    )
+
+    @model_validator(mode="after")
+    def _check_shapes(self) -> "Collated":
+        if (self.input_ids.shape != self.token_mask.shape) or (
+            self.conditions.shape != self.condition_mask.shape
+        ):
+            raise ValueError("Mask shapes don't match source")
+        if (self.input_ids.dim() != 2) or (self.conditions.dim() != 2):
+            raise ValueError("Tensors should have a batch dimension")
+        if (self.token_mask.dtype != torch.bool) or (
+            self.condition_mask.dtype != torch.bool
+        ):
+            raise ValueError("Mask tensors should have bool dtype")
+        return self
+
+
+class Collator[T](ABC):
+    """Base class for collators.
+
+    Provides a specification for preparing various data types for model input.
+    """
+
+    def __init__(self, model_config: ModelConfig, tokenizer: Tokenizer):
+        """Instantiates a Collator object.
+
+        Args:
+            model_config: Model configuration object.
+            tokenizer: Object implementing the Tokenizer protocol.
+        """
+        self.model_config = model_config
+        """Model configuration object."""
+        self.tokenizer = tokenizer
+        """Object implementing the Tokenizer protocol."""
+
+    def __call__(
+        self,
+        dataset: T,
+        seq_feature: str,
+        cond_features: list[str],
+        padding: bool,
+        add_special_tokens: bool,
+        device: torch.device | None,
+        **kwargs,
+    ) -> Collated:
+        return self.collate_data(
+            dataset,
+            seq_feature,
+            cond_features,
+            padding,
+            add_special_tokens,
+            device,
+            **kwargs,
+        )
+
+    @abstractmethod
+    def collate_data(
+        self,
+        dataset: T,
+        seq_feature: str,
+        cond_features: list[str],
+        padding: bool,
+        add_special_tokens: bool,
+        device: torch.device | None,
+        **kwargs,
+    ) -> Collated:
+        """Interface method for implementing data collation.
+
+        Args:
+            dataset: Dataset object of target data type.
+            seq_feature: Name of sequence feature or column.
+            cond_features: List of names of condition features or columns.
+            padding: Whether or not to uniformly pad outputs to the maximum length specified in the model configuration.
+            add_special_tokens: Whether or not to add special tokens according to a template. Special tokens will be
+                added by default during training, but will be omitted during inference, where the specified bos_token_id
+                will be appended before each pass through the decoder.
+            device: Torch device to which the produced tensors are moved.
+            **kwargs: Keyword arguments.
+
+        Returns:
+            Collated: Return schema representing a store of tensors ready for direct input to a Model.
+        """
+        ...
+
+
+class DataframeCollator(Collator):
+    """Collator designed for pandas.DataFrame inputs."""
+
+    def collate_data(
+        self,
+        dataset: DataFrame,
+        seq_feature: str,
+        cond_features: list[str],
+        padding: bool,
+        add_special_tokens: bool,
+        device: torch.device | None,
+        **kwargs,
+    ) -> Collated:
+        """Implementation of data collation for pandas.DataFrame inputs.
+
+        Args:
+            dataset: Dataset object of pandas.DataFrame type.
+            seq_feature: Name of sequence column.
+            cond_features: List of names of condition columns.
+            padding: Whether or not to uniformly pad outputs to the maximum length specified in the model configuration.
+            add_special_tokens: Whether or not to add special tokens according to a template. Special tokens will be
+                added by default during training, but will be omitted during inference, where the specified bos_token_id
+                will be appended before each pass through the decoder.
+            device: Torch device to which the produced tensors are moved.
+            **kwargs: Keyword arguments passed to Tokenizer.encode.
+
+        Returns:
+            Collated: Return schema representing a store of tensors ready for direct input to a Model.
+        """
+        df = dataset.copy()
+        input_ids = self.tokenizer.encode(
+            seq=df[seq_feature],
+            padding=padding,
+            max_length=self.model_config.input.max_len,
+            add_special_tokens=add_special_tokens,
+            **kwargs,
+        )
+        conditions = torch.from_numpy(df[cond_features].to_numpy()).to(
+            device=device, dtype=torch.float
+        )
+        token_mask = (input_ids != self.model_config.common.padding_idx).to(torch.bool)
+        condition_mask = (conditions != self.model_config.common.padding_value).to(
+            torch.bool
+        )
+        return Collated(
+            input_ids=input_ids,
+            conditions=conditions,
+            token_mask=token_mask,
+            condition_mask=condition_mask,
+        )

kernel_elastic_autoencoder-1.0.0/src/kernel_elastic_autoencoder/config.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import math
+from pathlib import Path
+
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    FilePath,
+    ImportString,
+    NonNegativeInt,
+    PositiveFloat,
+    PositiveInt,
+)
+
+
+class Config(BaseModel):
+    model_config = ConfigDict(frozen=True)
+
+    def to_dict(self) -> dict:
+        return self.model_dump()
+
+    @classmethod
+    def from_dict(cls, d: dict) -> Config:
+        return cls.model_validate(d)
+
+    def to_json(self, f: str):
+        return Path(f).write_text(self.model_dump_json(indent=2))
+
+    @classmethod
+    def from_json(cls, f: FilePath):
+        return cls.model_validate_json(Path(f).read_text())
+
+
+class ExperimentConfig(Config):
+    """Configuration schema for an experiment.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    model: ModelConfig = Field(..., description="ModelConfig object.")
+    training: TrainingConfig = Field(..., description="TrainingConfig object.")
+
+
+class ModelInputConfig(Config):
+    """ModelConfig schema block containing input parameters.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    max_len: PositiveInt = Field(
+        default=150,
+        description="Maximum sequence length through the encoder and decoder.",
+    )
+    vocab_size: PositiveInt = Field(
+        default=50,
+        description="Vocabulary size of tokenizer. Should be fetched with Tokenizer.vocab_size as specified in the "
+        "Tokenizer Protocol.",
+    )
+    condition_channels: PositiveInt = Field(
+        default=2,
+        description="Number of condition channels, should correspond to number of numerical columns in your "
+        "dataset/inputs.",
+    )
+
+
+class ModelCommonConfig(Config):
+    """ModelConfig schema block containing common architecture parameters.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    embedding_dim: PositiveInt = Field(
+        default=128, description="Embedding dimension used by nn.Embedding layers."
+    )
+    pooling_dim: PositiveInt = Field(
+        default=10,
+        description="Sequence length dimension to which inputs are pooled after condition concatenation through the "
+        "encoder. Proportional to the dimension of latent vectors.",
+    )
+    padding_idx: NonNegativeInt = Field(
+        default=0,
+        description="Index of padding token. Used internally to zero vectors corresponding to padding tokens through "
+        "embedding layers. Should be fetched with Tokenizer.pad_token_id as specified in the Tokenizer "
+        "protocol.",
+    )
+    padding_value: float = Field(
+        default=-1e10,
+        description='Value of "padded" condition fields. Should be set according to the value you use as a '
+        "placeholder for blank numerical fields in your dataset.",
+    )
+
+
+class ModelEncoderConfig(Config):
+    """ModelConfig schema block containing architecture parameters of the encoder.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    num_layers: PositiveInt = Field(
+        default=6, description="Number of Transformer encoder layers in the encoder."
+    )
+    num_heads: PositiveInt = Field(
+        default=4,
+        description="Number of self-attention heads used in each Transformer encoder layer in the encoder.",
+    )
+    feedforward_scale: PositiveInt = Field(
+        default=4,
+        description="Factor by which the dimension of the hidden layer in the FFNs differs from the dimension of the "
+        "input in the encoder. Applies to Transformer and Compression FFNs.",
+    )
+    dropout: PositiveFloat = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description="Dropout rate applied to Transformer encoder layers in the encoder.",
+    )
+
+
+class ModelDecoderConfig(Config):
+    """ModelConfig schema block containing architecture parameters of the decoder.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    num_layers: PositiveInt = Field(
+        default=6, description="Number of Transformer decoder layers in the decoder."
+    )
+    num_heads: PositiveInt = Field(
+        default=4,
+        description="Number of self- and cross-attention heads used in each Transformer decoder layer in the decoder.",
+    )
+    feedforward_scale: PositiveInt = Field(
+        default=4,
+        description="Factor by which the dimension of the hidden layer in the FFNs differs from the dimension of the "
+        "input in the decoder. Applies to Transformer and Mixing FFNs.",
+    )
+    dropout: PositiveFloat = Field(
+        default=0.1,
+        ge=0.0,
+        le=1.0,
+        description="Dropout rate applied to Transformer decoder layers in the decoder.",
+    )
+
+
+class ModelConfig(Config):
+    """Configuration schema for a Model. Contains all parameters needed to instantiate a Model.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    input: ModelInputConfig = Field(
+        default=ModelInputConfig(), description="ModelInputConfig object."
+    )
+    common: ModelCommonConfig = Field(
+        default=ModelCommonConfig(), description="ModelCommonConfig object."
+    )
+    encoder: ModelEncoderConfig = Field(
+        default=ModelEncoderConfig(), description="ModelEncoderConfig object."
+    )
+    decoder: ModelDecoderConfig = Field(
+        default=ModelDecoderConfig(), description="ModelDecoderConfig object."
+    )
+
+
+class TrainingCommonConfig(Config):
+    """TrainingConfig schema block containing common training parameters.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    max_epochs: PositiveInt = Field(
+        default=200, description="Maximum number of epochs to train for."
+    )
+    batch_size: PositiveInt = Field(
+        default=64,
+        description="Batch size for training. Large values are highly prone to OOM because batches are all padded to "
+        "a fixed length, as opposed to dynamically.",
+    )
+
+
+class TrainingHyperparameterConfig(Config):
+    """TrainingConfig schema block containing hyperparameters.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    hp_lambda: float = Field(
+        default=3.5,
+        description=r"Hyperparameter $\lambda$, as used in WCEL and m-MMD losses. Roughly, controls how strongly the "
+        r"shape of the latent vector distribution is penalized.",
+    )
+    hp_delta: float = Field(
+        default=1.0,
+        description=r"Hyperparameter $\delta$, as used in WCEL loss. Roughly, controls the relative weights of the "
+        r"vanilla-AE and VAE objectives in the reconstruction loss.",
+    )
+    hp_sigma: float = Field(
+        default=math.sqrt(32),
+        description=r"Hyperparameter $\sigma$, as used in the Kernel function applied in m-MMD loss. Roughly, "
+        r"used as a scaling factor to control the sizes of gradients produced by the m-MMD loss.",
+    )
+    kernel_dist_size: PositiveInt = Field(
+        default=1000,
+        description="Size of the sampled distribution of vectors used to penalize the shape of the latent vector "
+        "distribution through the kernel.",
+    )
+
+
+class TrainingOptimizerConfig(Config):
+    """TrainingConfig schema block containing configurations for the optimizer and scheduler.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    optimizer_fn: ImportString = Field(
+        default="torch.optim.AdamW",
+        description="Optimizer function import string. Should come from torch.optim.",
+    )
+    optimizer_params: dict = Field(
+        default={},
+        description="Dictionary of optimizer parameters to pass to **kwargs.",
+    )
+    scheduler_fn: ImportString = Field(
+        default="torch.optim.lr_scheduler.LinearLR",
+        description="Scheduler function import string. Should come from torch.optim.lr_scheduler.",
+    )
+    scheduler_params: dict = Field(
+        default={},
+        description="Dictionary of scheduler parameters to pass to **kwargs.",
+    )
+
+
+class TrainingConfig(Config):
+    """Configuration schema for a Trainer. Contains all parameters needed to instantiate a Trainer. Hyperparameter defaults are set according to experimentally-determined best practice in the original paper.
+
+    TODO: Pydantic configs will likely be deprecated or heavily refactored in the near future.
+    """
+
+    common: TrainingCommonConfig = Field(
+        default=TrainingCommonConfig(), description="TrainingCommonConfig object."
+    )
+    hyperparameters: TrainingHyperparameterConfig = Field(
+        default=TrainingHyperparameterConfig(),
+        description="TrainingHyperparameterConfig object.",
+    )
+    optimizer: TrainingOptimizerConfig = Field(
+        default=TrainingOptimizerConfig(), description="TrainingOptimizerConfig object."
+    )