kernel-elastic-autoencoder 3.3.5__tar.gz → 3.4.0__tar.gz

{kernel_elastic_autoencoder-3.3.5 → kernel_elastic_autoencoder-3.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 3.3.5
+Version: 3.4.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.3.5 → kernel_elastic_autoencoder-3.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.3.5"
+version = "3.4.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.3.5 → kernel_elastic_autoencoder-3.4.0}/src/kernel_elastic_autoencoder/losses.py

@@ -61,8 +61,11 @@ class Loss(nn.Module):
         """Sequence length dimension to which inputs are pooled after condition concatenation through the
                 encoder. Proportional to the dimension of latent vectors."""
 
-        self.register_buffer("_loc", torch.zeros(self.pooling_dim * self.embedding_dim))
-        self.register_buffer("_cov", torch.eye(self.pooling_dim * self.embedding_dim))
+        mvn = torch.distributions.MultivariateNormal(
+            loc=torch.zeros(self.pooling_dim * self.embedding_dim),
+            covariance_matrix=torch.eye(self.pooling_dim * self.embedding_dim),
+        )
+        self.register_buffer("_samples", mvn.sample((self.kernel_dist_size,)))
 
     def forward(
         self,
@@ -110,11 +113,7 @@ class Loss(nn.Module):
         self,
         latents: torch.Tensor,
     ) -> torch.Tensor:
-        mvn = torch.distributions.MultivariateNormal(
-            loc=self._loc,  # type: ignore
-            covariance_matrix=self._cov,  # type: ignore
-        )
-        samples = mvn.rsample((self.kernel_dist_size,)).to(latents.device)
+        samples = self._samples
         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.3.5 → kernel_elastic_autoencoder-3.4.0}/src/kernel_elastic_autoencoder/pipeline.py

@@ -52,6 +52,7 @@ class Pipeline:
         latents: torch.Tensor,
         sequences: list[str],
         conditions: list[list[float]] | torch.Tensor,
+        with_grad: bool = False,
         **kwargs,
     ) -> Iterable[str]:
         """Completes each conditioned input sequence.
@@ -108,7 +109,7 @@ class Pipeline:
                 .indices.squeeze(-1)
                 .to(torch.long)
             )
-            batches_completed |= new_toks.squeeze(-1) == self.tokenizer.eos_token_id
+            batches_completed |= new_toks.flatten() == self.tokenizer.eos_token_id
             new_toks = torch.where(
                 batches_completed.unsqueeze(-1),
                 self.tokenizer.pad_token_id,
@@ -118,9 +119,7 @@ class Pipeline:
             token_mask = torch.cat(
                 [
                     token_mask,
-                    torch.zeros(
-                        input_ids.size(0), 1, dtype=torch.bool, device=self.device
-                    ),
+                    batches_completed.unsqueeze(-1),
                 ],
                 dim=1,
             )
@@ -136,10 +135,9 @@ class Pipeline:
     ) -> Iterable[str]:
         """Completes each conditioned input sequence, using the beam search strategy.
 
-        The model completes each sequence in the provided list using decoder-only inference. On the
-        first step, the top `beam_size` first tokens are chosen for each batch. Subsequent tokens are
-        sampled greedily in parallel for every subsequence. Before returning output, the subsequence with
-        the highest sum of token logits is chosen for each batch.
+        The model completes each sequence in the provided list using decoder-only inference. The beam search
+        decoding algorithm is used. On the last step, instead of returning beam_size candidates per batch, the
+        candidate with the highest length-normalized sum of log odds is selected for each batch.
 
         Args:
             latents: Tensor of dimension (B, P * E) containing latent vectors for the batch.
@@ -184,14 +182,15 @@ class Pipeline:
             condition_embeddings=conds_embed,
             token_mask=token_mask,
         )
+        odds = logits.log_softmax(dim=-1)
         new_toks = (
-            torch.topk(logits[:, -1:], k=beam_size, dim=-1)
+            torch.topk(odds[:, -1:], k=beam_size, dim=-1)
             .indices.flatten()
             .unsqueeze(-1)
             .to(torch.long)
         )
         new_probs = (
-            torch.topk(logits[:, -1:], k=beam_size, dim=-1)
+            torch.topk(odds[:, -1:], k=beam_size, dim=-1)
             .values.flatten()
             .unsqueeze(-1)
         )
@@ -212,11 +211,11 @@ class Pipeline:
         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
+        batches_completed |= new_toks.flatten() == self.tokenizer.eos_token_id
         token_mask = torch.cat(
             [
                 token_mask,
-                torch.zeros(input_ids.size(0), 1, dtype=torch.bool, device=self.device),
+                batches_completed.unsqueeze(-1),
             ],
             dim=1,
         )
@@ -230,41 +229,81 @@ class Pipeline:
                 condition_embeddings=conds_embed,
                 token_mask=token_mask,
             )
+            odds = logits.log_softmax(dim=-1)
             new_toks = (
-                torch.topk(logits[:, -1:], k=1, dim=-1)
+                torch.topk(odds[:, -1:], k=beam_size, dim=-1)
                 .indices.flatten()
                 .unsqueeze(-1)
                 .to(torch.long)
             )
             new_probs = (
-                torch.topk(logits[:, -1:], k=1, dim=-1).values.flatten().unsqueeze(-1)
+                torch.topk(odds[:, -1:], k=beam_size, dim=-1)
+                .values.flatten()
+                .unsqueeze(-1)
             )
             new_toks = torch.where(
-                batches_completed.unsqueeze(-1),
+                batches_completed.repeat_interleave(beam_size, dim=0).unsqueeze(-1),
                 self.tokenizer.pad_token_id,
                 new_toks,
             )
             new_probs = torch.where(
-                batches_completed.unsqueeze(-1),
+                batches_completed.repeat_interleave(beam_size, dim=0).unsqueeze(-1),
                 0.0,
                 new_probs,
             )
-            input_ids = torch.cat([input_ids, new_toks], dim=1)
-            input_probs = torch.cat([input_probs, new_probs], dim=1)
+
+            candidate_ids = torch.cat(
+                [input_ids.repeat_interleave(beam_size, dim=0), new_toks], dim=1
+            )
+            candidate_probs = torch.cat(
+                [input_probs.repeat_interleave(beam_size, dim=0), new_probs], dim=1
+            )
+            top_probs = candidate_probs.view(
+                candidate_probs.size(0) // (beam_size**2), beam_size**2, -1
+            )
+            grouped_ids = candidate_ids.view(top_probs.size(0), top_probs.size(1), -1)
+            top_prob_inds = (
+                (
+                    top_probs.sum(dim=-1)
+                    / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
+                    .to(torch.long)
+                    .sum(dim=-1)
+                )
+                .topk(k=beam_size, dim=1)
+                .indices.squeeze(-1)
+            )
+            input_ids = grouped_ids[
+                torch.arange(top_probs.size(0)).unsqueeze(-1).repeat(1, beam_size),
+                top_prob_inds,
+            ].view(input_ids.size(0), -1)
+            input_probs = top_probs[
+                torch.arange(top_probs.size(0)).unsqueeze(-1).repeat(1, beam_size),
+                top_prob_inds,
+            ].view(input_ids.size(0), -1)
+
+            batches_completed |= (
+                input_ids[:, -1:].flatten() == self.tokenizer.eos_token_id
+            )
             token_mask = torch.cat(
                 [
                     token_mask,
-                    torch.zeros(
-                        input_ids.size(0), 1, dtype=torch.bool, device=self.device
-                    ),
+                    batches_completed.unsqueeze(-1),
                 ],
                 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.squeeze(-1)
         grouped_ids = input_ids.view(top_probs.shape[0], beam_size, -1)
+        top_prob_inds = (
+            (
+                top_probs.sum(dim=-1)
+                / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
+                .to(torch.long)
+                .sum(dim=-1)
+            )
+            .topk(k=1, dim=1)
+            .indices.squeeze(-1)
+        )
         winning_ids = grouped_ids[torch.arange(top_probs.size(0)), top_prob_inds]
         return self.tokenizer.decode(winning_ids, skip_special_tokens=True)
 

{kernel_elastic_autoencoder-3.3.5 → kernel_elastic_autoencoder-3.4.0}/src/kernel_elastic_autoencoder/training.py

@@ -1,5 +1,7 @@
+import json
 import os
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
+from typing import Any
 
 import torch
 from accelerate import Accelerator
@@ -37,6 +39,7 @@ class Trainer:
         conditions: Iterable[Iterable[float]] | torch.Tensor,
         train_split: float = 0.9,
         checkpoint: str = "./checkpoint",
+        epoch_callback: Callable[[dict[str, Any]], Any] = lambda _: None,
     ):
         """Trains a model. Optionally, resumes from an existing checkpoint.
 
@@ -49,6 +52,7 @@ class Trainer:
             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.
+            epoch_callback: Callback function accepting a dict of per-epoch stats.
         """
         accelerator = Accelerator(
             kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]
@@ -94,13 +98,15 @@ class Trainer:
             dataset_train,
             batch_size=self.config_typed.common.batch_size,
             pin_memory=True,
-            num_workers=0,
+            num_workers=4,
+            shuffle=True,
         )
         dataloader_test = torch.utils.data.DataLoader(
             dataset_test,
             batch_size=self.config_typed.common.batch_size,
             pin_memory=True,
-            num_workers=0,
+            num_workers=4,
+            shuffle=True,
         )
         curr_epoch = 0
 
@@ -121,7 +127,16 @@ class Trainer:
             curr_epoch = scheduler.scheduler.last_epoch + 1
 
         accelerator.wait_for_everyone()
+        all_contexts = []
         for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
+            if accelerator.is_main_process:
+                cb_ctx = {
+                    "epoch": epoch,
+                    "train_loss": None,
+                    "test_loss": None,
+                    "dist_mean": None,
+                    "dist_var": None,
+                }
             model.train()
             train_loss = torch.tensor([], device=accelerator.device)
             for input_ids, conditions, token_mask, condition_mask in tqdm(
@@ -131,16 +146,21 @@ class Trainer:
                 prediction, prediction_noise, latents = model(
                     input_ids, conditions, token_mask, condition_mask
                 )
-                loss = loss_fn(
-                    prediction, prediction_noise, input_ids[:, 1:], latents
-                )
+                loss = loss_fn(prediction, prediction_noise, input_ids[:, 1:], latents)
                 accelerator.backward(loss)
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
                 optimizer.step()
                 train_loss = torch.cat([train_loss, loss.detach().unsqueeze(-1)], dim=0)
             accelerator.print(f"Avg. train loss: {train_loss.mean().item()}")
+            if accelerator.is_main_process:
+                cb_ctx["train_loss"] = train_loss.mean().item()
 
             model.eval()
             test_loss = torch.tensor([], device=accelerator.device)
+
+            dist_mean = torch.tensor([], device=accelerator.device)
+            dist_var = torch.tensor([], device=accelerator.device)
             for input_ids, conditions, token_mask, condition_mask in tqdm(
                 dataloader_test, desc=f"Epoch {epoch}, Test Batch"
             ):
@@ -154,9 +174,22 @@ class Trainer:
                     loss = loss_fn(
                         prediction, prediction_noise, input_ids[:, 1:], latents
                     )
-                    test_loss = torch.cat([test_loss, loss.detach().unsqueeze(-1)], dim=0)
+                    mean = latents.mean().mean()
+                    var = latents.var(dim=0).mean()
+                    test_loss = torch.cat(
+                        [test_loss, loss.detach().unsqueeze(-1)], dim=0
+                    )
+                    dist_mean = torch.cat(
+                        [dist_mean, mean.detach().unsqueeze(-1)], dim=0
+                    )
+                    dist_var = torch.cat([dist_var, var.detach().unsqueeze(-1)], dim=0)
             accelerator.print(f"Avg. test loss: {test_loss.mean().item()}")
 
+            if accelerator.is_main_process:
+                cb_ctx["test_loss"] = test_loss.mean().item()
+                cb_ctx["dist_mean"] = dist_mean.mean().item()
+                cb_ctx["dist_var"] = dist_var.mean().item()
+
             scheduler.step()
 
             curr_epoch += 1
@@ -164,6 +197,10 @@ class Trainer:
             accelerator.save_state(checkpoint)
             if accelerator.is_main_process:
                 os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
+                all_contexts.append(cb_ctx)
+                with open("log.json", "w") as f:
+                    json.dump(all_contexts, f, indent=4)
+                epoch_callback(cb_ctx)
                 accelerator.unwrap_model(model).save_pretrained(
                     os.path.join(checkpoint, "dist/")
                 )