kernel-elastic-autoencoder 1.0.0__tar.gz → 3.0.0__tar.gz

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-3.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 1.0.0
+Version: 3.0.0
 Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
 License: MIT
 Author: Felix Rotter-McCartney
@@ -11,8 +11,8 @@ 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: accelerate (>=1.14.0,<2.0.0)
 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-1.0.0 → kernel_elastic_autoencoder-3.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "1.0.0"
+version = "3.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" }
@@ -11,13 +11,16 @@ 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)",
+    "accelerate (>=1.14.0,<2.0.0)",
 ]
 
 [tool.poetry]
 packages = [{ include = "kernel_elastic_autoencoder", from = "src" }]
 
+[tool.ruff.lint]
+ignore = ["B008"]
+
 [tool.semantic_release]
 version_toml = ["pyproject.toml:project.version"]
 commit_parser = "conventional"
@@ -40,7 +43,6 @@ dev = [
     "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 → kernel_elastic_autoencoder-3.0.0}/src/kernel_elastic_autoencoder/__init__.py

@@ -1,4 +1,3 @@
-from kernel_elastic_autoencoder.collate import Collated, Collator, DataframeCollator
 from kernel_elastic_autoencoder.config import (
     ExperimentConfig,
     ModelCommonConfig,
@@ -13,32 +12,24 @@ from kernel_elastic_autoencoder.config import (
 )
 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.pipeline import Pipeline
 from kernel_elastic_autoencoder.tokenizer import Tokenizer
-from kernel_elastic_autoencoder.training import Trainer, TrainerCallback
+from kernel_elastic_autoencoder.training import Trainer
 
 __all__ = [
     "ExperimentConfig",
-    "ModelConfig",
+    "Loss",
+    "Model",
     "ModelCommonConfig",
-    "ModelInputConfig",
-    "ModelEncoderConfig",
+    "ModelConfig",
     "ModelDecoderConfig",
-    "TrainingConfig",
-    "TrainingCommonConfig",
-    "TrainingHyperparameterConfig",
-    "TrainingOptimizerConfig",
+    "ModelEncoderConfig",
+    "ModelInputConfig",
     "Pipeline",
-    "Completion",
-    "Model",
-    "Loss",
     "Tokenizer",
-    "Collator",
-    "Collated",
-    "DataframeCollator",
-    "Sampler",
-    "Top1Sampler",
     "Trainer",
-    "TrainerCallback",
+    "TrainingCommonConfig",
+    "TrainingConfig",
+    "TrainingHyperparameterConfig",
+    "TrainingOptimizerConfig",
 ]

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-3.0.0}/src/kernel_elastic_autoencoder/losses.py

@@ -121,7 +121,7 @@ class Loss(nn.Module):
             / (2 * (self.hp_sigma**2))
         ).sum()
         loss = self.hp_lambda * (
-            torch.tensor([1])
+            torch.tensor(1, dtype=latents.dtype, device=latents.device)
             - ((1 / (latents.size(0) * self.kernel_dist_size)) * kernel_pairwise_sum)
         )
         return loss

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-3.0.0}/src/kernel_elastic_autoencoder/model.py

@@ -1,9 +1,5 @@
-from pathlib import Path
-from typing import Any
-
 import torch
 from huggingface_hub import PyTorchModelHubMixin
-from huggingface_hub.hub_mixin import T, DataclassInstance
 from torch import nn
 
 from kernel_elastic_autoencoder.config import ModelConfig
@@ -78,47 +74,6 @@ class Model(
         )
         """Model decoder module."""
 
-    @classmethod
-    def from_pretrained(
-        cls: type[T],
-        pretrained_model_name_or_path: str | Path,
-        *,
-        force_download: bool = False,
-        token: str | bool | None = None,
-        cache_dir: str | Path | None = None,
-        local_files_only: bool = False,
-        revision: str | None = None,
-        **model_kwargs,
-    ) -> T:
-        return cls.from_pretrained(
-            pretrained_model_name_or_path,
-            force_download=force_download,
-            token=token,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-            revision=revision,
-            **model_kwargs,
-        )
-
-    def save_pretrained(
-        self,
-        save_directory: str | Path,
-        *,
-        config: dict | DataclassInstance | None = None,
-        repo_id: str | None = None,
-        push_to_hub: bool = False,
-        model_card_kwargs: dict[str, Any] | None = None,
-        **push_to_hub_kwargs,
-    ) -> str | None:
-        return super().save_pretrained(
-            save_directory,
-            config=config,
-            repo_id=repo_id,
-            push_to_hub=push_to_hub,
-            model_card_kwargs=model_card_kwargs,
-            **push_to_hub_kwargs,
-        )
-
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -218,5 +173,13 @@ class Model(
         ).to(torch.bool)
         return self.decoder(current_output, latents, condition_embeddings, padding_mask)
 
-    def embed_conditions(self, conditions):
+    def embed_conditions(self, conditions: torch.Tensor) -> torch.Tensor:
+        """Basic interface for the separate embedding of condition vectors.
+
+        Args:
+            conditions: Tensor of dimension (B, C) containing condition values for each sequence.
+
+        Returns:
+            torch.Tensor: Tensor of dimension (B, C, E) containing condition embeddings for each sequence.
+        """
         return self.encoder.embedding.conditional_embedding(conditions)

kernel_elastic_autoencoder-3.0.0/src/kernel_elastic_autoencoder/pipeline.py

@@ -0,0 +1,161 @@
+from collections.abc import Iterable
+
+import torch
+
+from kernel_elastic_autoencoder.model import Model
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+
+
+class Pipeline:
+    """User-facing pipeline for inference.
+
+    Defines an easy-to-use API for decoder-only inference with a pretrained model.
+
+    Examples:
+        Wrapping a pretrained model and tokenizer:
+        >>> model = Model.from_pretrained("./checkpoint")
+        >>> tokenizer = MyTokenizer.from_pretrained("./checkpoint/tokenizer")
+        >>> pipe = Pipeline(model, tokenizer)
+
+        Getting a completion for sequences:
+        >>> compl = pipe.completion(latents, ["abc", "def", "ghi"], [[1.0, 0.5], [2.0, 1.0], [3.0, 1.5]])
+        >>> print(compl.outputs)
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        tokenizer: Tokenizer,
+        device: torch.device | None = None,
+    ) -> None:
+        """Instantiates a Pipeline object.
+
+        Args:
+            model: Pre-trained Model object for inference. Can be obtained from Model.from_pretrained.
+            tokenizer: Pre-configured Tokenizer object. The Tokenizer protocol supports tokenizers
+                inheriting from transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded
+                from HuggingFace Hub.
+            device: Torch device used for inference.
+        """
+        self.model = model.to(device)
+        """Pre-trained Model object for inference. Moved to Pipeline.device, and placed in eval() mode."""
+        self.model.eval()
+        self.tokenizer = tokenizer
+        """Pre-configured Tokenizer object."""
+        self.device = device
+        """Torch device used for inference."""
+
+    def _ingest(
+        self,
+        sequences: list[str],
+        conditions: list[list[float]] | torch.Tensor,
+        device: torch.device | None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        input_ids = self.tokenizer.encode(
+            seq=sequences,
+            padding=False,
+            max_length=self.model.config_typed.input.max_len,
+            add_special_tokens=False,
+            **kwargs,
+        )
+        conditions = torch.as_tensor(conditions, dtype=torch.float, device=device)
+        condition_mask = (
+            conditions != self.model.config_typed.common.padding_value
+        ).to(torch.bool)
+        return input_ids, conditions, condition_mask
+
+    @torch.inference_mode()
+    def _completion_entry(
+        self,
+        latents: torch.Tensor,
+        sequences: list[str],
+        conditions: list[list[float]] | torch.Tensor,
+        device: torch.device | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        input_ids, conds, cond_mask = self._ingest(
+            sequences, conditions, device, **kwargs
+        )
+        conds_embed = self.model.embed_conditions(conds)
+        return {
+            "latents": latents,
+            "input_ids": input_ids,
+            "batches_completed": torch.zeros(input_ids.size(0), dtype=torch.bool),
+            "condition_embeddings": conds_embed,
+            "condition_mask": cond_mask,
+        }
+
+    @torch.inference_mode()
+    def _completion_step(
+        self, intermediate: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        intermediate["input_ids"] = torch.cat(
+            [
+                torch.full(
+                    (intermediate["input_ids"].size(0), 1), self.tokenizer.bos_token_id
+                ),
+                intermediate["input_ids"],
+            ],
+            dim=1,
+        )
+        logits = self.model.decode(
+            current_output=intermediate["input_ids"],
+            latents=intermediate["latents"],
+            condition_embeddings=intermediate["condition_embeddings"],
+            token_mask=None,
+            condition_mask=intermediate["condition_mask"],
+        )
+        new_toks = (
+            torch.topk(logits[:, -1:], k=1, dim=-1).indices.squeeze(-1).to(torch.long)
+        )
+        intermediate["batches_completed"] |= (
+            new_toks.squeeze(-1) == self.tokenizer.eos_token_id
+        )
+        new_toks = torch.where(
+            intermediate["batches_completed"].unsqueeze(-1),
+            self.tokenizer.pad_token_id,
+            new_toks,
+        )
+        intermediate["input_ids"] = torch.cat([intermediate["input_ids"], new_toks], dim=1)
+        return intermediate
+
+    @torch.inference_mode()
+    def _completion_exit(self, intermediate: dict[str, torch.Tensor]) -> Iterable[str]:
+        return self.tokenizer.decode(intermediate["input_ids"], skip_special_tokens=True)
+
+    def completion(
+        self,
+        latents: torch.Tensor,
+        sequences: list[str],
+        conditions: list[list[float]] | torch.Tensor,
+        device: torch.device | None = None,
+        **kwargs,
+    ) -> Iterable[str]:
+        """Completes each conditioned input sequence.
+
+        The model completes each sequence in the provided list using decoder-only inference. Tokens are
+        sampled greedily.
+
+        Args:
+            latents: Tensor of dimension (B, P * E) containing latent vectors for the batch.
+            sequences: List of text sequences to complete.
+            conditions: List of condition value lists per batch.
+            device: Torch device used for inference.
+            **kwargs: Additional keyword arguments passed to Tokenizer.encode.
+
+        Returns:
+            Iterable[str]: List of completed sequences, stripped of special tokens.
+        """
+        intermediate = self._completion_entry(
+            latents=latents,
+            sequences=sequences,
+            conditions=conditions,
+            device=device,
+            **kwargs,
+        )
+        while (
+            intermediate["input_ids"].size(1) < self.model.config_typed.input.max_len
+        ) and (not intermediate["batches_completed"].all()):
+            intermediate = self._completion_step(intermediate=intermediate)
+        return self._completion_exit(intermediate=intermediate)

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-3.0.0}/src/kernel_elastic_autoencoder/tokenizer.py

@@ -78,71 +78,3 @@ class Tokenizer(Protocol):
             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-3.0.0/src/kernel_elastic_autoencoder/training.py

@@ -0,0 +1,153 @@
+import os
+from collections.abc import Iterable
+
+import torch
+from accelerate import Accelerator
+from accelerate.utils import tqdm
+
+from kernel_elastic_autoencoder.config import TrainingConfig
+from kernel_elastic_autoencoder.losses import Loss
+from kernel_elastic_autoencoder.model import Model
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+
+
+class Trainer:
+    def __init__(
+        self,
+        config: dict | TrainingConfig,
+    ):
+        """Instantiates a Trainer object.
+
+        Args:
+            config: Dictionary or TrainingConfig schema defining model parameters. Will be validated with TrainingConfig
+                regardless of input type.
+        """
+        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."""
+
+    def train(
+        self,
+        model: Model,
+        tokenizer: Tokenizer,
+        sequences: Iterable[str],
+        conditions: Iterable[Iterable[float]] | torch.Tensor,
+        train_split: float = 0.9,
+        checkpoint: str = "./checkpoint",
+    ):
+        """Trains a model. Optionally, resumes from an existing checkpoint.
+
+        Args:
+            model: Freshly instantiated model.
+            tokenizer: Pre-configured Tokenizer object. The Tokenizer protocol supports tokenizers
+                inheriting from transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded
+                from HuggingFace Hub.
+            sequences: Iterable of text sequences.
+            conditions: Iterable of iterables of condition values per sequence.
+            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.
+        """
+        accelerator = Accelerator()
+
+        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,
+        )
+        optimizer = self.config_typed.optimizer.optimizer_fn(
+            model.parameters(), **self.config_typed.optimizer.optimizer_params
+        )
+        scheduler = self.config_typed.optimizer.scheduler_fn(
+            optimizer, **self.config_typed.optimizer.scheduler_params
+        )
+
+        input_ids = tokenizer.encode(
+            seq=sequences,
+            padding=True,
+            max_length=model.config_typed.input.max_len,
+            add_special_tokens=True,
+        )
+        conditions = torch.as_tensor(conditions, dtype=torch.float)
+        token_mask = (input_ids != model.config_typed.common.padding_idx).to(torch.bool)
+        condition_mask = (conditions != model.config_typed.common.padding_value).to(
+            torch.bool
+        )
+
+        dataset = torch.utils.data.TensorDataset(
+            input_ids, conditions, token_mask, condition_mask
+        )
+        dataset_train, dataset_test = torch.utils.data.random_split(
+            dataset, [train_split, 1 - train_split]
+        )
+        dataloader_train = torch.utils.data.DataLoader(
+            dataset_train,
+            batch_size=self.config_typed.common.batch_size,
+        )
+        dataloader_test = torch.utils.data.DataLoader(
+            dataset_test,
+            batch_size=self.config_typed.common.batch_size,
+        )
+        curr_epoch = 0
+
+        model, optimizer, dataloader_train, dataloader_test, scheduler, curr_epoch = (
+            accelerator.prepare(
+                model,
+                optimizer,
+                dataloader_train,
+                dataloader_test,
+                scheduler,
+                curr_epoch,
+            )
+        )
+        accelerator.register_for_checkpointing(scheduler, curr_epoch)
+
+        if os.path.exists(checkpoint):
+            accelerator.load_state(checkpoint)
+
+        for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
+            model.train()
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_train, desc=f"Epoch {epoch}, Train Batch"
+            ):
+                optimizer.zero_grad()
+                prediction, prediction_noise, latents_noise = model(
+                    input_ids, conditions, token_mask, condition_mask
+                )
+                loss = loss_fn(
+                    prediction, prediction_noise, input_ids[:, 1:], latents_noise
+                )
+                accelerator.backward(loss)
+                optimizer.step()
+
+            model.eval()
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_test, desc=f"Epoch {epoch}, Test Batch"
+            ):
+                with torch.no_grad():
+                    prediction, prediction_noise, latents_noise = model(
+                        input_ids,
+                        conditions,
+                        token_mask,
+                        condition_mask,
+                    )
+                    loss = loss_fn(
+                        prediction, prediction_noise, input_ids[:, 1:], latents_noise
+                    )
+
+            scheduler.step(epoch)
+
+            accelerator.wait_for_everyone()
+            curr_epoch += 1
+            os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
+            model.save_pretrained(os.path.join(checkpoint, "dist/"))
+            self.config_typed.to_json(
+                os.path.join(checkpoint, "dist/train_config.json")
+            )
+            accelerator.save_state(checkpoint)