kernel-elastic-autoencoder 1.0.0__py3-none-any.whl

kernel_elastic_autoencoder/tokenizer.py

@@ -0,0 +1,148 @@
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Protocol, Self
+
+import torch
+
+
+class Tokenizer(Protocol):
+    """Protocol to be implemented for tokenizers. Supports transformers.PreTrainedTokenizerBase.
+
+    Provides a specification for tokenizing text for model input, and recovering text from token indices.
+    """
+
+    vocab_size: int
+    """Number of tokens in the tokenizer's vocabulary."""
+    bos_token: str
+    """Token marking the beginning of a sequence."""
+    bos_token_id: int
+    """Index of token marking the beginning of a sequence."""
+    eos_token: str
+    """Token marking the end of a sequence."""
+    eos_token_id: int
+    """Index of token marking the end of a sequence."""
+    pad_token: str
+    """Token representing padding."""
+    pad_token_id: int
+    """Index of token representing padding."""
+
+    def encode(
+        self,
+        seq: Iterable[str],
+        padding: bool,
+        max_length: int,
+        add_special_tokens: bool,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Encodes an Iterable of sequences to a tensor of indices. Optionally, adds special tokens according to a
+        template and pads the outputs to a fixed length.
+
+        Args:
+            seq: Iterable of text sequences to encode.
+            padding: Whether to pad the sequences to a fixed length.
+            max_length: Maximum length to which sequences are padded if padding is True.
+            add_special_tokens: Whether to add special tokens according to a template.
+            **kwargs: Keyword arguments.
+
+        Returns:
+            torch.Tensor: Tensor of dimension (B, S) containing vocabulary indices.
+        """
+        ...
+
+    def decode(
+        self, ids: torch.Tensor, skip_special_tokens: bool, **kwargs
+    ) -> Iterable[str]:
+        """Decodes a tensor of indices to an Iterable of sequences. Optionally, skips special tokens.
+
+        Args:
+            ids: Tensor of dimension (B, S) containing vocabulary indices to decode.
+            skip_special_tokens: Whether to skip decoding special tokens when constructing outputs.
+            **kwargs: Keyword arguments.
+
+        Returns:
+            Iterable[str]: Iterable of decoded sequences.
+        """
+        ...
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_model_name_or_path: str | Path, **kwargs
+    ) -> Self:
+        """Loads a tokenizer from a path or remote name. Supports Hugging Face from_pretrained.
+
+        Args:
+            pretrained_model_name_or_path: Local path or remote name.
+            **kwargs: Keyword arguments.
+
+        Returns:
+            Tokenizer: Pretrained tokenizer.
+        """
+        ...
+
+
+class _DummySlowTokenizer(Tokenizer):
+    def __init__(self, train: Iterable[str]):
+        self.pad_token = "_"
+        self.bos_token = "?"
+        self.eos_token = "!"
+
+        vocab = set[str]()
+        for b in train:
+            vocab = vocab.union(set(b))
+
+        special_tokens = [self.pad_token, self.bos_token, self.eos_token]
+        special_tokens.extend(list(vocab))
+        self.vocab = special_tokens
+        self.vocab_size = len(self.vocab)
+
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+
+    def encode(
+        self,
+        seq: Iterable[str],
+        padding: bool,
+        max_length: int,
+        add_special_tokens: bool,
+        **kwargs,
+    ) -> torch.Tensor:
+        ids_col = list[torch.Tensor]()
+        for b in seq:
+            toks = list(b)
+            if add_special_tokens:
+                toks.insert(0, self.bos_token)
+                toks.append(self.eos_token)
+            if padding:
+                toks += [self.pad_token] * (max_length - len(toks))
+            ids = torch.tensor([self.vocab.index(t) for t in toks], dtype=torch.long)
+            ids_col.append(ids)
+        return torch.stack(ids_col, dim=0)
+
+    def decode(
+        self, ids: torch.Tensor, skip_special_tokens: bool, **kwargs
+    ) -> Iterable[str]:
+        if skip_special_tokens:
+            return [
+                str(
+                    [
+                        (
+                            self.vocab[idx]
+                            if idx
+                            not in [
+                                self.pad_token_id,
+                                self.bos_token_id,
+                                self.eos_token_id,
+                            ]
+                            else ""
+                        )
+                        for idx in b
+                    ]
+                )
+                for b in ids
+            ]
+        else:
+            return [str([self.vocab[idx] for idx in b]) for b in ids]
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): ...

kernel_elastic_autoencoder/training.py

@@ -0,0 +1,386 @@
+import os
+from typing import Protocol
+
+import torch
+import torch.distributed as dist
+from pydantic import BaseModel, ConfigDict
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DistributedSampler
+
+from kernel_elastic_autoencoder.collate import Collated
+from kernel_elastic_autoencoder.config import TrainingConfig
+from kernel_elastic_autoencoder.losses import Loss
+from kernel_elastic_autoencoder.model import Model
+
+
+def _is_local():
+    return (
+        dist.is_torchelastic_launched() and (dist.get_rank() == 0)
+    ) or not dist.is_torchelastic_launched()
+
+
+class TrainerCallbackCtx(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    dist: bool
+    device_type: str
+    local_rank: int | None
+    model: Model | DDP
+    optimizer: torch.optim.Optimizer
+    scheduler: torch.optim.lr_scheduler.LRScheduler
+    dataloader_train: torch.utils.data.DataLoader
+    dataloader_test: torch.utils.data.DataLoader
+
+    epoch: int | None = None
+    rel_epoch: int | None = None
+    batch: int | None = None
+    train_loss: torch.Tensor | None = None
+    test_loss: torch.Tensor | None = None
+
+
+class TrainerCallback(Protocol):
+    """Protocol to be implemented for callback classes to a Trainer.
+
+    On each hook, except before the Trainer is initialized, a TrainerCallbackCtx schema is passed with the appropriate
+    information passed.
+    """
+
+    def before_init(self):
+        """Called before training setup."""
+        ...
+
+    def after_init(self, ctx: TrainerCallbackCtx):
+        """Called after training setup."""
+        ...
+
+    def before_epoch(self, ctx: TrainerCallbackCtx):
+        """Called before each epoch."""
+        ...
+
+    def before_train_batch(self, ctx: TrainerCallbackCtx):
+        """Called before each forward pass of a single training batch."""
+        ...
+
+    def after_train_batch(self, ctx: TrainerCallbackCtx):
+        """Called after each forward pass of a single training batch."""
+        ...
+
+    def before_test_batch(self, ctx: TrainerCallbackCtx):
+        """Called before each forward pass of a single test batch."""
+        ...
+
+    def after_test_batch(self, ctx: TrainerCallbackCtx):
+        """Called after each forward pass of a single test batch."""
+        ...
+
+    def after_epoch(self, ctx: TrainerCallbackCtx):
+        """Called after each epoch."""
+        ...
+
+    def after_training(self, ctx: TrainerCallbackCtx):
+        """Called after training ends."""
+        ...
+
+
+class TrainerDefaultCallback:
+    # TODO: Implement sensible default logging.
+    def before_init(self):
+        pass
+
+    def after_init(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def before_epoch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def before_train_batch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def after_train_batch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def before_test_batch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def after_test_batch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def after_epoch(self, ctx: TrainerCallbackCtx):
+        pass
+
+    def after_training(self, ctx: TrainerCallbackCtx):
+        pass
+
+
+class Trainer:
+    def __init__(
+        self,
+        config: dict | TrainingConfig,
+        callbacks: tuple[TrainerCallback] = (TrainerDefaultCallback(),),
+    ):
+        """Instantiates a Trainer object.
+
+        Args:
+            config: Dictionary or TrainingConfig schema defining model parameters. Will be validated with TrainingConfig
+                regardless of input type.
+            callbacks: Tuple of callback classes implementing TrainerCallback.
+        """
+        self.config = config
+        """Configuration object for Hugging Face Hub compatible serialization. Not recommended to use, as it is 
+                internal-use. Use Trainer.config_typed instead."""
+        config_typed = TrainingConfig.model_validate(config)
+        self.config_typed: TrainingConfig = config_typed
+        """Type-validated config in a Pydantic TrainingConfig schema, recommended for public API use."""
+        self.callbacks = callbacks
+        """Tuple of callback classes implementing TrainerCallback."""
+
+        self._ctx: TrainerCallbackCtx
+
+        self._model: DDP | Model
+        self._dist: bool
+        self._loss_fn: nn.Module
+        self._optimizer: torch.optim.Optimizer
+        self._scheduler: torch.optim.lr_scheduler.LRScheduler
+        self._dataloader_train: torch.utils.data.DataLoader
+        self._dataloader_test: torch.utils.data.DataLoader
+        self._device_type: str
+        self._local_rank: int | None
+
+    def train(
+        self,
+        model: Model,
+        ds: Collated,
+        train_split: float = 0.9,
+        checkpoint: str = "./checkpoint",
+    ):
+        """Trains a model with a Collated dataset. Optionally resumes from an existing checkpoint.
+
+        Args:
+            model: Freshly instantiated model.
+            ds: Tensor dataset following the Collated schema.
+            train_split: Fraction of dataset used for training. Must be between 0 and 1.
+            checkpoint: Path of local checkpoint to be saved and/or resumed.
+        """
+        next_epoch = self._setup(model, ds, train_split, checkpoint)
+        for epoch in range(next_epoch, self.config_typed.common.max_epochs):
+            self._ctx.rel_epoch = epoch - next_epoch
+            self._epoch(epoch, checkpoint)
+        self._ctx.epoch = None
+        self._ctx.rel_epoch = None
+        if _is_local():
+            for cb in self.callbacks:
+                cb.after_training(self._ctx)
+
+    def _setup(
+        self,
+        model: Model,
+        ds: Collated,
+        train_split: float,
+        checkpoint: str,
+    ):
+        if _is_local():
+            for cb in self.callbacks:
+                cb.before_init()
+        self._model = model
+        self._dist = dist.is_torchelastic_launched()
+        if self._dist:
+            self._model, self._device_type, self._local_rank = self._setup_ddp()
+        else:
+            self._device_type, self._local_rank = self._setup_no_ddp()
+        self._loss_fn = self._setup_loss(model)
+        self._optimizer, self._scheduler = self._setup_optimizer()
+        self._dataloader_train, self._dataloader_test = self._setup_dataloaders(
+            ds, train_split
+        )
+        next_epoch = 0
+        if os.path.exists(checkpoint):
+            next_epoch = self._resume(checkpoint)
+        self._ctx = TrainerCallbackCtx(
+            dist=self._dist,
+            device_type=self._device_type,
+            local_rank=self._local_rank,
+            model=self._model,
+            optimizer=self._optimizer,
+            scheduler=self._scheduler,
+            dataloader_train=self._dataloader_train,
+            dataloader_test=self._dataloader_test,
+        )
+        if _is_local():
+            for cb in self.callbacks:
+                cb.after_init(self._ctx)
+        return next_epoch
+
+    def _setup_loss(self, model: Model):
+        _loss_fn = Loss(
+            hp_lambda=self.config_typed.hyperparameters.hp_lambda,
+            hp_delta=self.config_typed.hyperparameters.hp_delta,
+            hp_sigma=self.config_typed.hyperparameters.hp_sigma,
+            kernel_dist_size=self.config_typed.hyperparameters.kernel_dist_size,
+            padding_idx=model.config_typed.common.padding_idx,
+            embedding_dim=model.config_typed.common.embedding_dim,
+            pooling_dim=model.config_typed.common.pooling_dim,
+        )
+        return _loss_fn
+
+    def _setup_optimizer(self):
+        optimizer = self.config_typed.optimizer.optimizer_fn(
+            self._model.parameters(), **self.config_typed.optimizer.optimizer_params
+        )
+        scheduler = self.config_typed.optimizer.scheduler_fn(
+            optimizer, **self.config_typed.optimizer.scheduler_params
+        )
+        return optimizer, scheduler
+
+    def _setup_ddp(self) -> tuple[DDP, str, int]:
+        device_type, vendor_backend = self._get_backend()
+        dist.init_process_group(backend=vendor_backend)
+        local_rank = int(os.environ["LOCAL_RANK"])
+        model = DDP(self._model.to(local_rank))
+        return model, device_type, local_rank
+
+    def _setup_no_ddp(self):
+        device_type, _ = self._get_backend()
+        return device_type, None
+
+    def _setup_dataloaders(self, ds: Collated, train_split: float):
+        dataset = torch.utils.data.TensorDataset(
+            ds.input_ids, ds.conditions, ds.token_mask, ds.condition_mask
+        )
+        dataset_train, dataset_test = torch.utils.data.random_split(
+            dataset, [train_split, 1 - train_split]
+        )
+        if self._dist:
+            sampler_train = DistributedSampler(dataset_train)
+            sampler_test = DistributedSampler(dataset_test)
+        else:
+            sampler_train = torch.utils.data.RandomSampler(dataset_train)
+            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
+        _dataloader_train = torch.utils.data.DataLoader(
+            dataset_train,
+            batch_size=self.config_typed.common.batch_size,
+            sampler=sampler_train,
+            pin_memory=True,
+        )
+        _dataloader_test = torch.utils.data.DataLoader(
+            dataset_test,
+            batch_size=self.config_typed.common.batch_size,
+            sampler=sampler_test,
+            pin_memory=True,
+        )
+        return _dataloader_train, _dataloader_test
+
+    def _resume(
+        self,
+        checkpoint: str,
+    ):
+        train_state = torch.load(os.path.join(checkpoint, "train_state.pt"))
+        self._optimizer.load_state_dict(train_state["optimizer"])
+        self._scheduler.load_state_dict(train_state["scheduler"])
+        self._model.from_pretrained(checkpoint)  # type: ignore
+        return train_state["next_epoch"]
+
+    def _epoch(
+        self,
+        epoch: int,
+        checkpoint: str,
+    ):
+        self._ctx.epoch = epoch
+        if _is_local():
+            for cb in self.callbacks:
+                cb.before_epoch(self._ctx)
+
+        self._model.train()
+        self._ctx.batch = 0
+        for input_ids, conditions, token_mask, condition_mask in self._dataloader_train:
+            if _is_local():
+                for cb in self.callbacks:
+                    cb.before_train_batch(self._ctx)
+            self._ctx.train_loss = self._batch_train(
+                input_ids, conditions, token_mask, condition_mask
+            )
+            if _is_local():
+                for cb in self.callbacks:
+                    cb.after_train_batch(self._ctx)
+            self._ctx.batch += 1
+
+        self._model.eval()
+        self._ctx.batch = 0
+        for input_ids, conditions, token_mask, condition_mask in self._dataloader_test:
+            if _is_local():
+                for cb in self.callbacks:
+                    cb.before_test_batch(self._ctx)
+            self._ctx.test_loss = self._batch_test(
+                input_ids, conditions, token_mask, condition_mask
+            )
+            if _is_local():
+                for cb in self.callbacks:
+                    cb.after_test_batch(self._ctx)
+            self._ctx.batch += 1
+
+        self._end_epoch(epoch, checkpoint)
+        self._ctx.train_loss = None
+        self._ctx.test_loss = None
+        self._ctx.batch = None
+        if _is_local():
+            for cb in self.callbacks:
+                cb.after_epoch(self._ctx)
+
+    def _end_epoch(self, epoch: int, checkpoint: str):
+        self._scheduler.step(epoch)
+        self._save_state(checkpoint, epoch)
+
+    def _batch_train(self, input_ids, conditions, token_mask, condition_mask):
+        self._optimizer.zero_grad()
+        with torch.amp.autocast(self._device_type):
+            prediction, prediction_noise, latents_noise = self._model.forward(
+                input_ids, conditions, token_mask, condition_mask
+            )
+            loss = self._loss_fn(
+                prediction, prediction_noise, input_ids[:, 1:], latents_noise
+            )
+        loss.backward()
+        self._optimizer.step()
+        return loss.detach()
+
+    def _batch_test(self, input_ids, conditions, token_mask, condition_mask):
+        with torch.no_grad():
+            with torch.amp.autocast(self._device_type):
+                prediction, prediction_noise, latents_noise = self._model.forward(
+                    input_ids,
+                    conditions,
+                    token_mask,
+                    condition_mask,
+                )
+            loss = self._loss_fn(
+                prediction, prediction_noise, input_ids[:, 1:], latents_noise
+            )
+        return loss.detach()
+
+    def _save_state(self, checkpoint: str, epoch: int):
+        if _is_local():
+            os.makedirs(checkpoint, exist_ok=True)
+            self._model.save_pretrained(checkpoint)  # type: ignore
+            torch.save(
+                {
+                    "next_epoch": epoch + 1,
+                    "optimizer": self._optimizer.state_dict(),
+                    "scheduler": self._scheduler.state_dict(),
+                },
+                os.path.join(checkpoint, "train_state.pt"),
+            )
+            self.config_typed.to_json(os.path.join(checkpoint, "train_config.json"))
+
+    def _get_backend(self) -> tuple[str, str]:
+        if torch.accelerator.is_available():
+            device_type = torch.accelerator.current_accelerator().type  # type: ignore
+            vendor_backend = torch.distributed.get_default_backend_for_device(
+                device_type
+            )
+
+        else:
+            device_type = torch.device("cpu").type
+            vendor_backend = torch.distributed.get_default_backend_for_device(
+                device_type
+            )
+
+        return device_type, vendor_backend

kernel_elastic_autoencoder-1.0.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,13 @@
+kernel_elastic_autoencoder/__init__.py,sha256=z7Vo43n0hcgYHHmPv2VTq9wvmRgD475F2Oa5cdysytQ,1227
+kernel_elastic_autoencoder/collate.py,sha256=k0p5Eryr7g7HplK6V8Ur9Gdg1sUR298zFWd8F-hGhzM,6114
+kernel_elastic_autoencoder/config.py,sha256=pm0i40_SS57Ur2cUpmJ1VoM20q-HY2D2veG8rULITm4,9052
+kernel_elastic_autoencoder/layers.py,sha256=QvZkyxij5x674E4YyhDDKIzsTAmT8-URlApWW2Vm2UU,6823
+kernel_elastic_autoencoder/losses.py,sha256=pMRKg-OmnuixLEz1yvmQw-EhYIGHU3wvVxGte4KY1Yo,6212
+kernel_elastic_autoencoder/model.py,sha256=30kLX1tPhaEycF0UoCsqE9EABSQ7_6aAwQO_U4Z2uTg,9128
+kernel_elastic_autoencoder/pipeline.py,sha256=39F294yrdBYPObPQp0r-rx2FlgPGEFZz8KiyHm_BUKI,8529
+kernel_elastic_autoencoder/sample.py,sha256=kRx4fZzQ169vdjWbtY-21yoNNX_3YmegIKxWYFhPU8U,1802
+kernel_elastic_autoencoder/tokenizer.py,sha256=-Fx-pV8s5rDxsZMbDKxwhz4Xmu8xUhIyAPtCE1Gii_E,4780
+kernel_elastic_autoencoder/training.py,sha256=FDpk059D94kYkfsi5l-1yzTjHI5U8Hm6PAP1FKTfGVQ,13742
+kernel_elastic_autoencoder-1.0.0.dist-info/METADATA,sha256=wkNaUu-lbFlxOmVqHfeabTCrfMg0iLr7QWnOh7tx_lY,2001
+kernel_elastic_autoencoder-1.0.0.dist-info/WHEEL,sha256=eY7nduwzv-ldUxpzbRlxwvC693Hg6PX8bWDjEHjZ_dk,88
+kernel_elastic_autoencoder-1.0.0.dist-info/RECORD,,

kernel_elastic_autoencoder-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.4.1
+Root-Is-Purelib: true
+Tag: py3-none-any