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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "3.3.6"
+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.6 → 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.6 → 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.
@@ -181,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)
         )
@@ -263,7 +265,7 @@ class Pipeline:
             top_prob_inds = (
                 (
                     top_probs.sum(dim=-1)
-                    * (grouped_ids != self.tokenizer.eos_token_id)
+                    / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
                     .to(torch.long)
                     .sum(dim=-1)
                 )
@@ -295,7 +297,7 @@ class Pipeline:
         top_prob_inds = (
             (
                 top_probs.sum(dim=-1)
-                * (grouped_ids != self.tokenizer.eos_token_id)
+                / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
                 .to(torch.long)
                 .sum(dim=-1)
             )

{kernel_elastic_autoencoder-3.3.6 → 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)]
@@ -123,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(
@@ -133,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"
             ):
@@ -156,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
@@ -166,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/")
                 )