kernel-elastic-autoencoder 3.1.3__tar.gz → 3.2.0__tar.gz

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.1.3
+Version: 3.2.0
 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.3 → kernel_elastic_autoencoder-3.2.0}/pyproject.toml

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

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/config.py

@@ -9,8 +9,8 @@ from pydantic import (
     Field,
     FilePath,
     ImportString,
+    NonNegativeFloat,
     NonNegativeInt,
-    PositiveFloat,
     PositiveInt,
 )
 
@@ -110,7 +110,7 @@ class ModelEncoderConfig(Config):
         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(
+    dropout: NonNegativeFloat = Field(
         default=0.0,
         ge=0.0,
         le=1.0,
@@ -136,7 +136,7 @@ class ModelDecoderConfig(Config):
         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(
+    dropout: NonNegativeFloat = Field(
         default=0.1,
         ge=0.0,
         le=1.0,

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/layers.py

@@ -31,7 +31,7 @@ class ConditionEmbedding(nn.Module):
 
     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, 0)
+        masked_indices = indices.masked_fill(c == self.padding_value, 0).to(c.device)
         embedding = self.embedding(masked_indices)
         return embedding * c.unsqueeze(-1).repeat(1, 1, embedding.size(-1))
 
@@ -205,7 +205,7 @@ class Decoder(nn.Module):
         self.out_linear = nn.Linear(embedding_dim, vocab_size)
 
         self.register_buffer(
-            "causal_mask", nn.Transformer.generate_square_subsequent_mask(max_len - 1).to(torch.bool)
+            "causal_mask", nn.Transformer.generate_square_subsequent_mask(max_len).to(torch.bool)
         )
 
     def forward(

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/losses.py

@@ -95,7 +95,7 @@ class Loss(nn.Module):
         prediction_noise: torch.Tensor,
         ground_truth: torch.Tensor,
     ) -> torch.Tensor:
-        ground_truth.to(torch.long)
+        ground_truth = ground_truth.to(torch.long)
         log_softmax = torch.log_softmax(prediction, dim=-1)
         log_softmax_noise = torch.log_softmax(prediction_noise, dim=-1)
         Y = F.one_hot(ground_truth, num_classes=prediction.size(-1))

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/model.py

@@ -108,13 +108,13 @@ class Model(
             input_ids[:, :-1],
             latents,
             condition_embeddings,
-            padding_mask[: self.config_typed.input.max_len - 1],
+            padding_mask[:, : self.config_typed.input.max_len - 1],
         )
         prediction_noise = self.decoder(
             input_ids[:, :-1],
             latents_noise,
             condition_embeddings,
-            padding_mask[: self.config_typed.input.max_len - 1],
+            padding_mask[:, : self.config_typed.input.max_len - 1],
         )
         return prediction, prediction_noise, latents_noise
 
@@ -139,8 +139,8 @@ class Model(
         padding_mask = (
             torch.cat((token_mask, condition_mask), 1)
             if (token_mask is not None)
-            else torch.cat((torch.full_like(input_ids, True), condition_mask), 1)
-        )
+            else torch.cat((torch.full_like(input_ids, False), condition_mask), 1)
+        ).to(torch.bool)
         return self.encoder(input_ids, conditions, padding_mask)[0]
 
     def decode(
@@ -149,7 +149,6 @@ class Model(
         latents: torch.Tensor,
         condition_embeddings: torch.Tensor,
         token_mask: torch.Tensor | None,
-        condition_mask: torch.Tensor,
     ) -> torch.Tensor:
         """Basic interface for a forward pass through the Model.model.decoder module.
 
@@ -161,15 +160,13 @@ class Model(
                 be produced using Model.embed_conditions.
             token_mask: Tensor of dimension (B, S) containing boolean padding masks for each sequence. If None is
                 passed, the absence of padding is assumed.
-            condition_mask: Tensor of dimension (B, C) containing boolean condition padding masks for each sequence.
 
         Returns:
             torch.Tensor: Tensor of dimension (B, S, L) containing prediction logits produced by the decoder.
         """
-        padding_mask = (
-            torch.cat((token_mask, condition_mask), 1)
+        padding_mask = (token_mask
             if (token_mask is not None)
-            else torch.cat((torch.full_like(current_output, False), condition_mask), 1)
+            else torch.full_like(current_output, False)
         ).to(torch.bool)
         return self.decoder(current_output, latents, condition_embeddings, padding_mask)
 

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/pipeline.py

@@ -37,20 +37,21 @@ class Pipeline:
                 from HuggingFace Hub.
             device: Torch device used for inference.
         """
-        self.model = model.to(device)
+        self.device = device
+        """Torch device used for inference."""
+        if self.device is None:
+            self.device = torch.device("cpu")
+        self.model = model.to(self.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 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.
@@ -62,7 +63,6 @@ class Pipeline:
             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:
@@ -70,38 +70,38 @@ class Pipeline:
         """
         input_ids = self.tokenizer.encode(
             text=sequences,
-            padding='do_not_pad',
-            max_length=self.model.config_typed.input.max_len,
+            padding="longest",
             add_special_tokens=False,
             return_tensors="pt",
             **kwargs,
         ).to(self.device)
-        conditions = torch.as_tensor(conditions, dtype=torch.float, device=device)
-        condition_mask = (
-            conditions != self.model.config_typed.common.padding_value
-        ).to(torch.bool)
+        input_ids = torch.cat(
+            [
+                torch.full(
+                    (input_ids.size(0), 1),
+                    self.tokenizer.bos_token_id,
+                    device=self.device,
+                ),
+                input_ids,
+            ],
+            dim=1,
+        )
+        latents = latents.to(self.device)
+        conditions = torch.as_tensor(conditions, dtype=torch.float, device=self.device)
         conds_embed = self.model.embed_conditions(conditions)
-        batches_completed = torch.zeros(input_ids.size(0), dtype=torch.bool)
+        token_mask = (input_ids == self.tokenizer.pad_token_id).to(torch.bool)
+        batches_completed = torch.zeros(
+            input_ids.size(0), dtype=torch.bool, device=self.device
+        )
 
         while (input_ids.size(1) < self.model.config_typed.input.max_len) and (
             not batches_completed.all()
         ):
-            input_ids = torch.cat(
-                [
-                    torch.full(
-                        (input_ids.size(0), 1),
-                        self.tokenizer.bos_token_id,
-                    ),
-                    input_ids,
-                ],
-                dim=1,
-            )
             logits = self.model.decode(
                 current_output=input_ids,
                 latents=latents,
                 condition_embeddings=conds_embed,
-                token_mask=None,
-                condition_mask=condition_mask,
+                token_mask=token_mask,
             )
             new_toks = (
                 torch.topk(logits[:, -1:], k=1, dim=-1)
@@ -115,6 +115,15 @@ class Pipeline:
                 new_toks,
             )
             input_ids = torch.cat([input_ids, new_toks], dim=1)
+            token_mask = torch.cat(
+                [
+                    token_mask,
+                    torch.zeros(
+                        input_ids.size(0), 1, dtype=torch.bool, device=self.device
+                    ),
+                ],
+                dim=1,
+            )
         return self.tokenizer.decode(input_ids, skip_special_tokens=True)
 
     def beam_completion(
@@ -123,7 +132,6 @@ class Pipeline:
         beam_size: int,
         sequences: list[str],
         conditions: list[list[float]] | torch.Tensor,
-        device: torch.device | None = None,
         **kwargs,
     ) -> Iterable[str]:
         """Completes each conditioned input sequence, using the beam search strategy.
@@ -138,7 +146,6 @@ class Pipeline:
             beam_size: Beam size of first step.
             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:
@@ -146,42 +153,48 @@ class Pipeline:
         """
         input_ids = self.tokenizer.encode(
             text=sequences,
-            padding='do_not_pad',
-            max_length=self.model.config_typed.input.max_len,
+            padding="longest",
             add_special_tokens=False,
             return_tensors="pt",
             **kwargs,
         ).to(self.device)
-        input_probs = torch.zeros_like(input_ids * beam_size)
-        conditions = torch.as_tensor(conditions, dtype=torch.float, device=device)
-        condition_mask = (
-            (conditions != self.model.config_typed.common.padding_value)
-            .to(torch.bool)
-            .repeat_interleave(beam_size, dim=0)
+        input_ids = torch.cat(
+            [
+                torch.full(
+                    (input_ids.size(0), 1),
+                    self.tokenizer.bos_token_id,
+                    device=self.device,
+                ),
+                input_ids,
+            ],
+            dim=1,
         )
-        conds_embed = self.model.embed_conditions(conditions).repeat_interleave(
-            beam_size, dim=0
+        token_mask = (input_ids == self.tokenizer.pad_token_id).to(torch.bool)
+        latents = latents.to(self.device)
+        input_probs = torch.zeros_like(input_ids).repeat_interleave(beam_size, dim=0)
+        conditions = torch.as_tensor(conditions, dtype=torch.float, device=self.device)
+        conds_embed = self.model.embed_conditions(conditions)
+        batches_completed = torch.zeros(
+            input_ids.size(0) * beam_size, dtype=torch.bool, device=self.device
         )
-        batches_completed = torch.zeros(input_ids.size(0) * beam_size, dtype=torch.bool)
 
         logits = self.model.decode(
             current_output=input_ids,
             latents=latents,
             condition_embeddings=conds_embed,
-            token_mask=None,
-            condition_mask=condition_mask,
+            token_mask=token_mask,
         )
         new_toks = (
             torch.topk(logits[:, -1:], k=beam_size, dim=-1)
-            .indices.squeeze(-1)
+            .indices.flatten()
+            .unsqueeze(-1)
             .to(torch.long)
         )
         new_probs = (
             torch.topk(logits[:, -1:], k=beam_size, dim=-1)
-            .values.squeeze(-1)
-            .to(torch.long)
+            .values.flatten()
+            .unsqueeze(-1)
         )
-        batches_completed |= new_toks.squeeze(-1) == self.tokenizer.eos_token_id
         new_toks = torch.where(
             batches_completed.unsqueeze(-1),
             self.tokenizer.pad_token_id,
@@ -196,38 +209,36 @@ class Pipeline:
             [input_ids.repeat_interleave(beam_size, dim=0), new_toks], dim=1
         )
         input_probs = torch.cat([input_probs, new_probs], dim=1)
+        conds_embed = conds_embed.repeat_interleave(beam_size, dim=0)
+        latents = latents.repeat_interleave(beam_size, dim=0)
+        token_mask = token_mask.repeat_interleave(beam_size, dim=0)
+        batches_completed |= new_toks.squeeze(-1) == self.tokenizer.eos_token_id
+        token_mask = torch.cat(
+            [
+                token_mask,
+                torch.zeros(input_ids.size(0), 1, dtype=torch.bool, device=self.device),
+            ],
+            dim=1,
+        )
 
         while (input_ids.size(1) < self.model.config_typed.input.max_len) and (
             not batches_completed.all()
         ):
-            input_ids = torch.cat(
-                [
-                    torch.full(
-                        (input_ids.size(0), 1),
-                        self.tokenizer.bos_token_id,
-                    ),
-                    input_ids,
-                ],
-                dim=1,
-            )
             logits = self.model.decode(
                 current_output=input_ids,
                 latents=latents,
                 condition_embeddings=conds_embed,
-                token_mask=None,
-                condition_mask=condition_mask,
+                token_mask=token_mask,
             )
             new_toks = (
                 torch.topk(logits[:, -1:], k=1, dim=-1)
-                .indices.squeeze(-1)
+                .indices.flatten()
+                .unsqueeze(-1)
                 .to(torch.long)
             )
             new_probs = (
-                torch.topk(logits[:, -1:], k=1, dim=-1)
-                .values.squeeze(-1)
-                .to(torch.long)
+                torch.topk(logits[:, -1:], k=1, dim=-1).values.flatten().unsqueeze(-1)
             )
-            batches_completed |= new_toks.squeeze(-1) == self.tokenizer.eos_token_id
             new_toks = torch.where(
                 batches_completed.unsqueeze(-1),
                 self.tokenizer.pad_token_id,
@@ -240,10 +251,44 @@ class Pipeline:
             )
             input_ids = torch.cat([input_ids, new_toks], dim=1)
             input_probs = torch.cat([input_probs, new_probs], dim=1)
+            token_mask = torch.cat(
+                [
+                    token_mask,
+                    torch.zeros(
+                        input_ids.size(0), 1, dtype=torch.bool, device=self.device
+                    ),
+                ],
+                dim=1,
+            )
+            batches_completed |= new_toks.squeeze(-1) == self.tokenizer.eos_token_id
 
         top_probs = input_probs.reshape(input_probs.size(0) // beam_size, beam_size, -1)
-        top_prob_inds = top_probs.sum(dim=-1).topk(k=1, dim=1).indices
-        winning_ids = torch.take_along_dim(
-            input_ids, top_prob_inds.unsqueeze(-1), 1
-        ).squeeze(1)
+        top_prob_inds = top_probs.sum(dim=-1).topk(k=1, dim=1).indices.squeeze(-1)
+        grouped_ids = input_ids.view(top_probs.shape[0], beam_size, -1)
+        winning_ids = grouped_ids[torch.arange(top_probs.size(0)), top_prob_inds]
         return self.tokenizer.decode(winning_ids, skip_special_tokens=True)
+
+    def encoding(
+        self, sequences: list[str], conditions: list[list[float]] | torch.Tensor
+    ) -> torch.Tensor:
+        input_ids = self.tokenizer.encode(
+            text=sequences,
+            padding="max_length",
+            max_length=self.model.config_typed.input.max_len,
+            add_special_tokens=True,
+            return_tensors="pt",
+        ).to(self.device)
+        conditions = torch.as_tensor(conditions, dtype=torch.float, device=self.device)
+        token_mask = (input_ids == self.model.config_typed.common.padding_idx).to(
+            dtype=torch.bool, device=self.device
+        )
+        condition_mask = (
+            conditions == self.model.config_typed.common.padding_value
+        ).to(dtype=torch.bool, device=self.device)
+
+        return self.model.encode(
+            input_ids=input_ids,
+            conditions=conditions,
+            token_mask=token_mask,
+            condition_mask=condition_mask,
+        )

{kernel_elastic_autoencoder-3.1.3 → kernel_elastic_autoencoder-3.2.0}/src/kernel_elastic_autoencoder/training.py

@@ -76,8 +76,8 @@ class Trainer:
             return_tensors="pt",
         )
         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(
+        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
         )
 
@@ -127,7 +127,8 @@ class Trainer:
                 )
                 accelerator.backward(loss)
                 optimizer.step()
-
+            print(f"Train loss: {float(loss.detach())}")
+            
             model.eval()
             for input_ids, conditions, token_mask, condition_mask in tqdm(
                 dataloader_test, desc=f"Epoch {epoch}, Test Batch"
@@ -142,8 +143,9 @@ class Trainer:
                     loss = loss_fn(
                         prediction, prediction_noise, input_ids[:, 1:], latents_noise
                     )
+            print(f"Test loss: {float(loss.detach())}")
 
-            scheduler.step(epoch)
+            scheduler.step()
 
             accelerator.wait_for_everyone()
             curr_epoch += 1