kernel-elastic-autoencoder 1.0.0__tar.gz → 2.0.0__tar.gz

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/PKG-INFO

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

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "1.0.0"
+version = "2.0.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" }
@@ -18,6 +18,9 @@ dependencies = [
 [tool.poetry]
 packages = [{ include = "kernel_elastic_autoencoder", from = "src" }]
 
+[tool.ruff.lint]
+ignore = ["B008"]
+
 [tool.semantic_release]
 version_toml = ["pyproject.toml:project.version"]
 commit_parser = "conventional"
@@ -40,7 +43,6 @@ dev = [
     "ruff (>=0.15.20,<0.17.0)",
     "pytest-cov (>=7.1.0,<8.0.0)",
     "pytest-xdist (>=3.8.0,<4.0.0)",
-    "ty (>=0.0.62,<0.0.64)",
     "python-semantic-release (>=10.6.1,<11.0.0)",
     "pdoc (>=16.0.0,<17.0.0)",
 ]

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/src/kernel_elastic_autoencoder/losses.py

@@ -121,7 +121,7 @@ class Loss(nn.Module):
             / (2 * (self.hp_sigma**2))
         ).sum()
         loss = self.hp_lambda * (
-            torch.tensor([1])
+            torch.tensor(1, dtype=latents.dtype, device=latents.device)
             - ((1 / (latents.size(0) * self.kernel_dist_size)) * kernel_pairwise_sum)
         )
         return loss

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/src/kernel_elastic_autoencoder/model.py

@@ -1,9 +1,6 @@
-from pathlib import Path
-from typing import Any
 
 import torch
 from huggingface_hub import PyTorchModelHubMixin
-from huggingface_hub.hub_mixin import T, DataclassInstance
 from torch import nn
 
 from kernel_elastic_autoencoder.config import ModelConfig
@@ -78,47 +75,6 @@ class Model(
         )
         """Model decoder module."""
 
-    @classmethod
-    def from_pretrained(
-        cls: type[T],
-        pretrained_model_name_or_path: str | Path,
-        *,
-        force_download: bool = False,
-        token: str | bool | None = None,
-        cache_dir: str | Path | None = None,
-        local_files_only: bool = False,
-        revision: str | None = None,
-        **model_kwargs,
-    ) -> T:
-        return cls.from_pretrained(
-            pretrained_model_name_or_path,
-            force_download=force_download,
-            token=token,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-            revision=revision,
-            **model_kwargs,
-        )
-
-    def save_pretrained(
-        self,
-        save_directory: str | Path,
-        *,
-        config: dict | DataclassInstance | None = None,
-        repo_id: str | None = None,
-        push_to_hub: bool = False,
-        model_card_kwargs: dict[str, Any] | None = None,
-        **push_to_hub_kwargs,
-    ) -> str | None:
-        return super().save_pretrained(
-            save_directory,
-            config=config,
-            repo_id=repo_id,
-            push_to_hub=push_to_hub,
-            model_card_kwargs=model_card_kwargs,
-            **push_to_hub_kwargs,
-        )
-
     def forward(
         self,
         input_ids: torch.Tensor,

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/src/kernel_elastic_autoencoder/pipeline.py

@@ -31,7 +31,7 @@ class Completion(BaseModel):
     """Tensor containing generated condition masks for reuse."""
 
 
-class Pipeline[_T]:
+class Pipeline[T]:
     """User-facing pipeline for inference.
 
     Defines an easy-to-use API for decoder-only inference with a pretrained model.
@@ -51,7 +51,7 @@ class Pipeline[_T]:
         self,
         model: Model,
         tokenizer: Tokenizer,
-        collator: Callable[[ModelConfig, Tokenizer], Collator[_T]] = DataframeCollator,
+        collator: Callable[[ModelConfig, Tokenizer], Collator[T]] = DataframeCollator,
         sampler: Callable[[Tokenizer], Sampler] = Top1Sampler,
         device: torch.device | None = None,
     ) -> None:
@@ -87,7 +87,7 @@ class Pipeline[_T]:
     def _completion_entry(
         self,
         latents: torch.Tensor,
-        dataset: _T,
+        dataset: T,
         seq_feature: str,
         cond_features: list[str],
         device: torch.device | None = None,
@@ -148,7 +148,7 @@ class Pipeline[_T]:
         self, intermediate: CompletionIntermediate, **kwargs
     ) -> Completion:
         return Completion(
-            outputs=self.tokenizer.decode(
+            outputs=self.sampler(
                 intermediate.input_ids, skip_special_tokens=True
             ),
             condition_embeddings=intermediate.condition_embeddings,
@@ -158,7 +158,7 @@ class Pipeline[_T]:
     def completion(
         self,
         latents: torch.Tensor,
-        dataset: _T,
+        dataset: T,
         seq_feature: str,
         cond_features: list[str],
         device: torch.device | None = None,

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/src/kernel_elastic_autoencoder/sample.py

@@ -12,25 +12,34 @@ class Sampler(ABC):
     Provides a specification for sampling text sequences from logits.
     """
 
-    def __init__(self, tokenizer: Tokenizer):
+    def __init__(self, tokenizer: Tokenizer, **kwargs):
         """Instantiates a Sampler object.
 
         Args:
             tokenizer: Object implementing the Tokenizer protocol.
+            **kwargs: Sampler-specific keyword args.
         """
         self.tokenizer = tokenizer
         """Object implementing the Tokenizer protocol."""
 
+    @abstractmethod
     def __call__(
         self, logits: torch.Tensor, skip_special_tokens: bool, **kwargs
     ) -> Iterable[str]:
-        ids = self.sample_ids(logits, **kwargs)
-        sampled = self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)
-        return sampled
+        """Interface method for implementing the last sampling step where outputs are decoded through the tokenizer.
+
+        Args:
+            logits: Tensor of dimension (B, S) containing IDs at the last step of inference.
+            **kwargs: Keyword arguments.
+
+        Returns:
+            Iterable[str]: Iterable of strings decoded by the tokenizer.
+        """
+        ...
 
     @abstractmethod
     def sample_ids(self, logits: torch.Tensor, **kwargs) -> torch.Tensor:
-        """Interface method for implementing index sampling from logits.
+        """Interface method for implementing index sampling from logits in intermediate steps.
 
         Args:
             logits: Tensor of dimension (B, S, L) containing logits.
@@ -43,6 +52,9 @@ class Sampler(ABC):
 
 
 class Top1Sampler(Sampler):
+    def __call__(self, logits: torch.Tensor, skip_special_tokens: bool, **kwargs) -> Iterable[str]:
+        return self.tokenizer.decode(logits, skip_special_tokens=skip_special_tokens)
+
     def sample_ids(self, logits: torch.Tensor, **kwargs) -> torch.Tensor:
         """Implementation of index sampling from logits choosing the highest-probability token.
 

{kernel_elastic_autoencoder-1.0.0 → kernel_elastic_autoencoder-2.0.0}/src/kernel_elastic_autoencoder/__init__.py

@@ -19,26 +19,26 @@ from kernel_elastic_autoencoder.tokenizer import Tokenizer
 from kernel_elastic_autoencoder.training import Trainer, TrainerCallback
 
 __all__ = [
+    "Collated",
+    "Collator",
+    "Completion",
+    "DataframeCollator",
     "ExperimentConfig",
-    "ModelConfig",
+    "Loss",
+    "Model",
     "ModelCommonConfig",
-    "ModelInputConfig",
-    "ModelEncoderConfig",
+    "ModelConfig",
     "ModelDecoderConfig",
-    "TrainingConfig",
-    "TrainingCommonConfig",
-    "TrainingHyperparameterConfig",
-    "TrainingOptimizerConfig",
+    "ModelEncoderConfig",
+    "ModelInputConfig",
     "Pipeline",
-    "Completion",
-    "Model",
-    "Loss",
-    "Tokenizer",
-    "Collator",
-    "Collated",
-    "DataframeCollator",
     "Sampler",
+    "Tokenizer",
     "Top1Sampler",
     "Trainer",
     "TrainerCallback",
+    "TrainingCommonConfig",
+    "TrainingConfig",
+    "TrainingHyperparameterConfig",
+    "TrainingOptimizerConfig",
 ]