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

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

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

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

@@ -52,7 +52,6 @@ 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.
@@ -131,6 +130,7 @@ class Pipeline:
         beam_size: int,
         sequences: list[str],
         conditions: list[list[float]] | torch.Tensor,
+        hp_alpha: float = 0.5,
         **kwargs,
     ) -> Iterable[str]:
         """Completes each conditioned input sequence, using the beam search strategy.
@@ -144,6 +144,7 @@ class Pipeline:
             beam_size: Beam size of first step.
             sequences: List of text sequences to complete.
             conditions: List of condition value lists per batch.
+            hp_alpha: Strength of length normalization.
             **kwargs: Additional keyword arguments passed to Tokenizer.encode.
 
         Returns:
@@ -265,9 +266,7 @@ class Pipeline:
             top_prob_inds = (
                 (
                     top_probs.sum(dim=-1)
-                    / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
-                    .to(torch.long)
-                    .sum(dim=-1)
+                    / torch.pow((grouped_ids != self.tokenizer.pad_token_id).to(torch.long).sum(dim=-1), hp_alpha)
                 )
                 .topk(k=beam_size, dim=1)
                 .indices.squeeze(-1)
@@ -297,9 +296,7 @@ class Pipeline:
         top_prob_inds = (
             (
                 top_probs.sum(dim=-1)
-                / torch.sqrt(grouped_ids != self.tokenizer.pad_token_id)
-                .to(torch.long)
-                .sum(dim=-1)
+                / torch.pow((grouped_ids != self.tokenizer.pad_token_id).to(torch.long).sum(dim=-1), hp_alpha)
             )
             .topk(k=1, dim=1)
             .indices.squeeze(-1)

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

@@ -194,7 +194,7 @@ class Trainer:
 
             curr_epoch += 1
             accelerator.wait_for_everyone()
-            accelerator.save_state(checkpoint)
+            accelerator.save_state(checkpoint, safe_serialization=False, save_on_each_node=True)
             if accelerator.is_main_process:
                 os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
                 all_contexts.append(cb_ctx)