kernel-elastic-autoencoder 3.1.0__tar.gz → 3.1.2__tar.gz

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.1.0
+Version: 3.1.2
 Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
 License: MIT
 Author: Felix Rotter-McCartney

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.1.0"
+version = "3.1.2"
 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" }

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/config.py

@@ -217,6 +217,7 @@ class TrainingOptimizerConfig(Config):
     optimizer_fn: ImportString = Field(
         default="torch.optim.AdamW",
         description="Optimizer function import string. Should come from torch.optim.",
+        validate_default=True,
     )
     optimizer_params: dict = Field(
         default={},
@@ -225,6 +226,7 @@ class TrainingOptimizerConfig(Config):
     scheduler_fn: ImportString = Field(
         default="torch.optim.lr_scheduler.LinearLR",
         description="Scheduler function import string. Should come from torch.optim.lr_scheduler.",
+        validate_default=True,
     )
     scheduler_params: dict = Field(
         default={},

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/layers.py

@@ -20,20 +20,18 @@ class ConditionEmbedding(nn.Module):
         self,
         condition_channels: int,
         embedding_dim: int,
-        padding_idx: int,
         padding_value: float,
     ):
         super().__init__()
         self.padding_value = padding_value
-        self.padding_idx = padding_idx
         self.embedding = nn.Embedding(
-            condition_channels + 1, embedding_dim, padding_idx=padding_idx
+            condition_channels + 1, embedding_dim, padding_idx=0
         )
         self.register_buffer("indices", torch.arange(1, condition_channels + 1))
 
     def forward(self, c: torch.Tensor):
         indices = self.indices.repeat(c.size(0), 1).to(c.device)  # type: ignore
-        masked_indices = indices.masked_fill(c == self.padding_value, self.padding_idx)
+        masked_indices = indices.masked_fill(c == self.padding_value, 0)
         embedding = self.embedding(masked_indices)
         return embedding * c.unsqueeze(-1).repeat(1, 1, embedding.size(-1))
 
@@ -56,7 +54,9 @@ class TransformerEmbedding(nn.Module):
             max_len, embedding_dim, padding_idx
         )
         self.conditional_embedding = ConditionEmbedding(
-            condition_channels, embedding_dim, padding_idx, padding_value
+            condition_channels,
+            embedding_dim,
+            padding_value,
         )
 
     def forward(self, x: torch.Tensor, c: torch.Tensor | None = None):

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/losses.py

@@ -114,7 +114,7 @@ class Loss(nn.Module):
             loc=self._loc,  # type: ignore
             covariance_matrix=self._cov,  # type: ignore
         )
-        samples = mvn.rsample((self.kernel_dist_size,))
+        samples = mvn.rsample((self.kernel_dist_size,)).to(latents.device)
         square_difference_sum = torch.cdist(latents, samples, p=2.0).pow(2)
         kernel_pairwise_sum = torch.exp(
             ((-1 / (self.pooling_dim * self.embedding_dim)) * square_difference_sum)

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/pipeline.py

@@ -69,10 +69,11 @@ class Pipeline:
             Iterable[str]: List of completed sequences, stripped of special tokens.
         """
         input_ids = self.tokenizer.encode(
-            seq=sequences,
-            padding=False,
+            text=sequences,
+            padding='do_not_pad',
             max_length=self.model.config_typed.input.max_len,
             add_special_tokens=False,
+            return_tensors="pt",
             **kwargs,
         ).to(self.device)
         conditions = torch.as_tensor(conditions, dtype=torch.float, device=device)
@@ -144,10 +145,11 @@ class Pipeline:
             Iterable[str]: List of completed sequences, stripped of special tokens.
         """
         input_ids = self.tokenizer.encode(
-            seq=sequences,
-            padding=False,
+            text=sequences,
+            padding='do_not_pad',
             max_length=self.model.config_typed.input.max_len,
             add_special_tokens=False,
+            return_tensors="pt",
             **kwargs,
         ).to(self.device)
         input_probs = torch.zeros_like(input_ids * beam_size)

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/tokenizer.py

@@ -11,8 +11,6 @@ class Tokenizer(Protocol):
     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
@@ -28,20 +26,22 @@ class Tokenizer(Protocol):
 
     def encode(
         self,
-        seq: Iterable[str],
-        padding: bool,
+        text: Iterable[str],
+        padding: str,
         max_length: int,
         add_special_tokens: bool,
+        return_tensors: str,
         **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.
+            text: 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.
+            return_tensors: Tensor return type, must be either 'pt' or 'np'.
             **kwargs: Keyword arguments.
 
         Returns:
@@ -50,12 +50,12 @@ class Tokenizer(Protocol):
         ...
 
     def decode(
-        self, ids: torch.Tensor, skip_special_tokens: bool, **kwargs
+        self, token_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.
+            token_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.
 
@@ -78,3 +78,5 @@ class Tokenizer(Protocol):
             Tokenizer: Pretrained tokenizer.
         """
         ...
+
+    def __len__(self) -> int: ...

{kernel_elastic_autoencoder-3.1.0 → kernel_elastic_autoencoder-3.1.2}/src/kernel_elastic_autoencoder/training.py

@@ -69,10 +69,11 @@ class Trainer:
         )
 
         input_ids = tokenizer.encode(
-            seq=sequences,
-            padding=True,
+            text=sequences,
+            padding="max_length",
             max_length=model.config_typed.input.max_len,
             add_special_tokens=True,
+            return_tensors="pt",
         )
         conditions = torch.as_tensor(conditions, dtype=torch.float)
         token_mask = (input_ids != model.config_typed.common.padding_idx).to(torch.bool)
@@ -106,10 +107,11 @@ class Trainer:
                 curr_epoch,
             )
         )
-        accelerator.register_for_checkpointing(scheduler, curr_epoch)
+        accelerator.register_for_checkpointing(scheduler)
 
         if os.path.exists(checkpoint):
             accelerator.load_state(checkpoint)
+            curr_epoch = scheduler.scheduler.last_epoch + 1
 
         for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
             model.train()