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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kernel_elastic_autoencoder
-Version: 2.0.0
+Version: 3.0.0
 Summary: Implementation of Kernel-Elastic Autoencoder for Molecular Design (https://doi.org/10.1093/pnasnexus/pgae168)
 License: MIT
 Author: Felix Rotter-McCartney
@@ -11,8 +11,8 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Requires-Dist: accelerate (>=1.14.0,<2.0.0)
 Requires-Dist: huggingface-hub (>=1.22.0,<2.0.0)
-Requires-Dist: pandas (>=3.0.3,<4.0.0)
 Requires-Dist: pydantic (>=2.13.4,<3.0.0)
 Requires-Dist: safetensors (>=0.8.0,<0.9.0)
 Description-Content-Type: text/markdown

{kernel_elastic_autoencoder-2.0.0 → kernel_elastic_autoencoder-3.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "kernel_elastic_autoencoder"
-version = "2.0.0"
+version = "3.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" }
@@ -11,8 +11,8 @@ requires-python = ">=3.12,<3.15"
 dependencies = [
     "pydantic (>=2.13.4,<3.0.0)",
     "huggingface-hub (>=1.22.0,<2.0.0)",
-    "pandas (>=3.0.3,<4.0.0)",
     "safetensors (>=0.8.0,<0.9.0)",
+    "accelerate (>=1.14.0,<2.0.0)",
 ]
 
 [tool.poetry]

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

@@ -1,4 +1,3 @@
-from kernel_elastic_autoencoder.collate import Collated, Collator, DataframeCollator
 from kernel_elastic_autoencoder.config import (
     ExperimentConfig,
     ModelCommonConfig,
@@ -13,16 +12,11 @@ from kernel_elastic_autoencoder.config import (
 )
 from kernel_elastic_autoencoder.losses import Loss
 from kernel_elastic_autoencoder.model import Model
-from kernel_elastic_autoencoder.pipeline import Completion, Pipeline
-from kernel_elastic_autoencoder.sample import Sampler, Top1Sampler
+from kernel_elastic_autoencoder.pipeline import Pipeline
 from kernel_elastic_autoencoder.tokenizer import Tokenizer
-from kernel_elastic_autoencoder.training import Trainer, TrainerCallback
+from kernel_elastic_autoencoder.training import Trainer
 
 __all__ = [
-    "Collated",
-    "Collator",
-    "Completion",
-    "DataframeCollator",
     "ExperimentConfig",
     "Loss",
     "Model",
@@ -32,11 +26,8 @@ __all__ = [
     "ModelEncoderConfig",
     "ModelInputConfig",
     "Pipeline",
-    "Sampler",
     "Tokenizer",
-    "Top1Sampler",
     "Trainer",
-    "TrainerCallback",
     "TrainingCommonConfig",
     "TrainingConfig",
     "TrainingHyperparameterConfig",

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

@@ -1,4 +1,3 @@
-
 import torch
 from huggingface_hub import PyTorchModelHubMixin
 from torch import nn
@@ -174,5 +173,13 @@ class Model(
         ).to(torch.bool)
         return self.decoder(current_output, latents, condition_embeddings, padding_mask)
 
-    def embed_conditions(self, conditions):
+    def embed_conditions(self, conditions: torch.Tensor) -> torch.Tensor:
+        """Basic interface for the separate embedding of condition vectors.
+
+        Args:
+            conditions: Tensor of dimension (B, C) containing condition values for each sequence.
+
+        Returns:
+            torch.Tensor: Tensor of dimension (B, C, E) containing condition embeddings for each sequence.
+        """
         return self.encoder.embedding.conditional_embedding(conditions)

kernel_elastic_autoencoder-3.0.0/src/kernel_elastic_autoencoder/pipeline.py

@@ -0,0 +1,161 @@
+from collections.abc import Iterable
+
+import torch
+
+from kernel_elastic_autoencoder.model import Model
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+
+
+class Pipeline:
+    """User-facing pipeline for inference.
+
+    Defines an easy-to-use API for decoder-only inference with a pretrained model.
+
+    Examples:
+        Wrapping a pretrained model and tokenizer:
+        >>> model = Model.from_pretrained("./checkpoint")
+        >>> tokenizer = MyTokenizer.from_pretrained("./checkpoint/tokenizer")
+        >>> pipe = Pipeline(model, tokenizer)
+
+        Getting a completion for sequences:
+        >>> compl = pipe.completion(latents, ["abc", "def", "ghi"], [[1.0, 0.5], [2.0, 1.0], [3.0, 1.5]])
+        >>> print(compl.outputs)
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        tokenizer: Tokenizer,
+        device: torch.device | None = None,
+    ) -> None:
+        """Instantiates a Pipeline object.
+
+        Args:
+            model: Pre-trained Model object for inference. Can be obtained from Model.from_pretrained.
+            tokenizer: Pre-configured Tokenizer object. The Tokenizer protocol supports tokenizers
+                inheriting from transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded
+                from HuggingFace Hub.
+            device: Torch device used for inference.
+        """
+        self.model = model.to(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 _ingest(
+        self,
+        sequences: list[str],
+        conditions: list[list[float]] | torch.Tensor,
+        device: torch.device | None,
+        **kwargs,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        input_ids = self.tokenizer.encode(
+            seq=sequences,
+            padding=False,
+            max_length=self.model.config_typed.input.max_len,
+            add_special_tokens=False,
+            **kwargs,
+        )
+        conditions = torch.as_tensor(conditions, dtype=torch.float, device=device)
+        condition_mask = (
+            conditions != self.model.config_typed.common.padding_value
+        ).to(torch.bool)
+        return input_ids, conditions, condition_mask
+
+    @torch.inference_mode()
+    def _completion_entry(
+        self,
+        latents: torch.Tensor,
+        sequences: list[str],
+        conditions: list[list[float]] | torch.Tensor,
+        device: torch.device | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        input_ids, conds, cond_mask = self._ingest(
+            sequences, conditions, device, **kwargs
+        )
+        conds_embed = self.model.embed_conditions(conds)
+        return {
+            "latents": latents,
+            "input_ids": input_ids,
+            "batches_completed": torch.zeros(input_ids.size(0), dtype=torch.bool),
+            "condition_embeddings": conds_embed,
+            "condition_mask": cond_mask,
+        }
+
+    @torch.inference_mode()
+    def _completion_step(
+        self, intermediate: dict[str, torch.Tensor]
+    ) -> dict[str, torch.Tensor]:
+        intermediate["input_ids"] = torch.cat(
+            [
+                torch.full(
+                    (intermediate["input_ids"].size(0), 1), self.tokenizer.bos_token_id
+                ),
+                intermediate["input_ids"],
+            ],
+            dim=1,
+        )
+        logits = self.model.decode(
+            current_output=intermediate["input_ids"],
+            latents=intermediate["latents"],
+            condition_embeddings=intermediate["condition_embeddings"],
+            token_mask=None,
+            condition_mask=intermediate["condition_mask"],
+        )
+        new_toks = (
+            torch.topk(logits[:, -1:], k=1, dim=-1).indices.squeeze(-1).to(torch.long)
+        )
+        intermediate["batches_completed"] |= (
+            new_toks.squeeze(-1) == self.tokenizer.eos_token_id
+        )
+        new_toks = torch.where(
+            intermediate["batches_completed"].unsqueeze(-1),
+            self.tokenizer.pad_token_id,
+            new_toks,
+        )
+        intermediate["input_ids"] = torch.cat([intermediate["input_ids"], new_toks], dim=1)
+        return intermediate
+
+    @torch.inference_mode()
+    def _completion_exit(self, intermediate: dict[str, torch.Tensor]) -> Iterable[str]:
+        return self.tokenizer.decode(intermediate["input_ids"], skip_special_tokens=True)
+
+    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.
+
+        The model completes each sequence in the provided list using decoder-only inference. Tokens are
+        sampled greedily.
+
+        Args:
+            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:
+            Iterable[str]: List of completed sequences, stripped of special tokens.
+        """
+        intermediate = self._completion_entry(
+            latents=latents,
+            sequences=sequences,
+            conditions=conditions,
+            device=device,
+            **kwargs,
+        )
+        while (
+            intermediate["input_ids"].size(1) < self.model.config_typed.input.max_len
+        ) and (not intermediate["batches_completed"].all()):
+            intermediate = self._completion_step(intermediate=intermediate)
+        return self._completion_exit(intermediate=intermediate)

{kernel_elastic_autoencoder-2.0.0 → kernel_elastic_autoencoder-3.0.0}/src/kernel_elastic_autoencoder/tokenizer.py

@@ -78,71 +78,3 @@ class Tokenizer(Protocol):
             Tokenizer: Pretrained tokenizer.
         """
         ...
-
-
-class _DummySlowTokenizer(Tokenizer):
-    def __init__(self, train: Iterable[str]):
-        self.pad_token = "_"
-        self.bos_token = "?"
-        self.eos_token = "!"
-
-        vocab = set[str]()
-        for b in train:
-            vocab = vocab.union(set(b))
-
-        special_tokens = [self.pad_token, self.bos_token, self.eos_token]
-        special_tokens.extend(list(vocab))
-        self.vocab = special_tokens
-        self.vocab_size = len(self.vocab)
-
-        self.pad_token_id = 0
-        self.bos_token_id = 1
-        self.eos_token_id = 2
-
-    def encode(
-        self,
-        seq: Iterable[str],
-        padding: bool,
-        max_length: int,
-        add_special_tokens: bool,
-        **kwargs,
-    ) -> torch.Tensor:
-        ids_col = list[torch.Tensor]()
-        for b in seq:
-            toks = list(b)
-            if add_special_tokens:
-                toks.insert(0, self.bos_token)
-                toks.append(self.eos_token)
-            if padding:
-                toks += [self.pad_token] * (max_length - len(toks))
-            ids = torch.tensor([self.vocab.index(t) for t in toks], dtype=torch.long)
-            ids_col.append(ids)
-        return torch.stack(ids_col, dim=0)
-
-    def decode(
-        self, ids: torch.Tensor, skip_special_tokens: bool, **kwargs
-    ) -> Iterable[str]:
-        if skip_special_tokens:
-            return [
-                str(
-                    [
-                        (
-                            self.vocab[idx]
-                            if idx
-                            not in [
-                                self.pad_token_id,
-                                self.bos_token_id,
-                                self.eos_token_id,
-                            ]
-                            else ""
-                        )
-                        for idx in b
-                    ]
-                )
-                for b in ids
-            ]
-        else:
-            return [str([self.vocab[idx] for idx in b]) for b in ids]
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): ...

kernel_elastic_autoencoder-3.0.0/src/kernel_elastic_autoencoder/training.py

@@ -0,0 +1,153 @@
+import os
+from collections.abc import Iterable
+
+import torch
+from accelerate import Accelerator
+from accelerate.utils import tqdm
+
+from kernel_elastic_autoencoder.config import TrainingConfig
+from kernel_elastic_autoencoder.losses import Loss
+from kernel_elastic_autoencoder.model import Model
+from kernel_elastic_autoencoder.tokenizer import Tokenizer
+
+
+class Trainer:
+    def __init__(
+        self,
+        config: dict | TrainingConfig,
+    ):
+        """Instantiates a Trainer object.
+
+        Args:
+            config: Dictionary or TrainingConfig schema defining model parameters. Will be validated with TrainingConfig
+                regardless of input type.
+        """
+        self.config = config
+        """Configuration object for Hugging Face Hub compatible serialization. Not recommended to use, as it is 
+                internal-use. Use Trainer.config_typed instead."""
+        config_typed = TrainingConfig.model_validate(config)
+        self.config_typed: TrainingConfig = config_typed
+        """Type-validated config in a Pydantic TrainingConfig schema, recommended for public API use."""
+
+    def train(
+        self,
+        model: Model,
+        tokenizer: Tokenizer,
+        sequences: Iterable[str],
+        conditions: Iterable[Iterable[float]] | torch.Tensor,
+        train_split: float = 0.9,
+        checkpoint: str = "./checkpoint",
+    ):
+        """Trains a model. Optionally, resumes from an existing checkpoint.
+
+        Args:
+            model: Freshly instantiated model.
+            tokenizer: Pre-configured Tokenizer object. The Tokenizer protocol supports tokenizers
+                inheriting from transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded
+                from HuggingFace Hub.
+            sequences: Iterable of text sequences.
+            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.
+        """
+        accelerator = Accelerator()
+
+        loss_fn = Loss(
+            hp_lambda=self.config_typed.hyperparameters.hp_lambda,
+            hp_delta=self.config_typed.hyperparameters.hp_delta,
+            hp_sigma=self.config_typed.hyperparameters.hp_sigma,
+            kernel_dist_size=self.config_typed.hyperparameters.kernel_dist_size,
+            padding_idx=model.config_typed.common.padding_idx,
+            embedding_dim=model.config_typed.common.embedding_dim,
+            pooling_dim=model.config_typed.common.pooling_dim,
+        )
+        optimizer = self.config_typed.optimizer.optimizer_fn(
+            model.parameters(), **self.config_typed.optimizer.optimizer_params
+        )
+        scheduler = self.config_typed.optimizer.scheduler_fn(
+            optimizer, **self.config_typed.optimizer.scheduler_params
+        )
+
+        input_ids = tokenizer.encode(
+            seq=sequences,
+            padding=True,
+            max_length=model.config_typed.input.max_len,
+            add_special_tokens=True,
+        )
+        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(
+            torch.bool
+        )
+
+        dataset = torch.utils.data.TensorDataset(
+            input_ids, conditions, token_mask, condition_mask
+        )
+        dataset_train, dataset_test = torch.utils.data.random_split(
+            dataset, [train_split, 1 - train_split]
+        )
+        dataloader_train = torch.utils.data.DataLoader(
+            dataset_train,
+            batch_size=self.config_typed.common.batch_size,
+        )
+        dataloader_test = torch.utils.data.DataLoader(
+            dataset_test,
+            batch_size=self.config_typed.common.batch_size,
+        )
+        curr_epoch = 0
+
+        model, optimizer, dataloader_train, dataloader_test, scheduler, curr_epoch = (
+            accelerator.prepare(
+                model,
+                optimizer,
+                dataloader_train,
+                dataloader_test,
+                scheduler,
+                curr_epoch,
+            )
+        )
+        accelerator.register_for_checkpointing(scheduler, curr_epoch)
+
+        if os.path.exists(checkpoint):
+            accelerator.load_state(checkpoint)
+
+        for epoch in range(curr_epoch, self.config_typed.common.max_epochs):
+            model.train()
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_train, desc=f"Epoch {epoch}, Train Batch"
+            ):
+                optimizer.zero_grad()
+                prediction, prediction_noise, latents_noise = model(
+                    input_ids, conditions, token_mask, condition_mask
+                )
+                loss = loss_fn(
+                    prediction, prediction_noise, input_ids[:, 1:], latents_noise
+                )
+                accelerator.backward(loss)
+                optimizer.step()
+
+            model.eval()
+            for input_ids, conditions, token_mask, condition_mask in tqdm(
+                dataloader_test, desc=f"Epoch {epoch}, Test Batch"
+            ):
+                with torch.no_grad():
+                    prediction, prediction_noise, latents_noise = model(
+                        input_ids,
+                        conditions,
+                        token_mask,
+                        condition_mask,
+                    )
+                    loss = loss_fn(
+                        prediction, prediction_noise, input_ids[:, 1:], latents_noise
+                    )
+
+            scheduler.step(epoch)
+
+            accelerator.wait_for_everyone()
+            curr_epoch += 1
+            os.makedirs(os.path.join(checkpoint, "dist/"), exist_ok=True)
+            model.save_pretrained(os.path.join(checkpoint, "dist/"))
+            self.config_typed.to_json(
+                os.path.join(checkpoint, "dist/train_config.json")
+            )
+            accelerator.save_state(checkpoint)

kernel_elastic_autoencoder-2.0.0/src/kernel_elastic_autoencoder/collate.py

@@ -1,164 +0,0 @@
-from abc import ABC, abstractmethod
-
-import torch
-from pandas import DataFrame
-from pydantic import BaseModel, ConfigDict, Field, model_validator
-
-from kernel_elastic_autoencoder.config import ModelConfig
-from kernel_elastic_autoencoder.tokenizer import Tokenizer
-
-
-class Collated(BaseModel):
-    """Return schema for Collator.__call__. Represents a store of tensors ready for direct input to a Model."""
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    input_ids: torch.Tensor = Field(
-        ...,
-        description="Tensor of dimension (B, S) containing indices of input tokens, generated by a Tokenizer.",
-    )
-    conditions: torch.Tensor = Field(
-        ...,
-        description="Tensor of dimension (B, C) containing condition values for each sequence.",
-    )
-    token_mask: torch.Tensor = Field(
-        ...,
-        description="Tensor of dimension (B, S) containing a boolean padding mask for each sequence.",
-    )
-    condition_mask: torch.Tensor = Field(
-        ...,
-        description="Tensor of dimension (B, C) containing a boolean condition mask for each sequence",
-    )
-
-    @model_validator(mode="after")
-    def _check_shapes(self) -> "Collated":
-        if (self.input_ids.shape != self.token_mask.shape) or (
-            self.conditions.shape != self.condition_mask.shape
-        ):
-            raise ValueError("Mask shapes don't match source")
-        if (self.input_ids.dim() != 2) or (self.conditions.dim() != 2):
-            raise ValueError("Tensors should have a batch dimension")
-        if (self.token_mask.dtype != torch.bool) or (
-            self.condition_mask.dtype != torch.bool
-        ):
-            raise ValueError("Mask tensors should have bool dtype")
-        return self
-
-
-class Collator[T](ABC):
-    """Base class for collators.
-
-    Provides a specification for preparing various data types for model input.
-    """
-
-    def __init__(self, model_config: ModelConfig, tokenizer: Tokenizer):
-        """Instantiates a Collator object.
-
-        Args:
-            model_config: Model configuration object.
-            tokenizer: Object implementing the Tokenizer protocol.
-        """
-        self.model_config = model_config
-        """Model configuration object."""
-        self.tokenizer = tokenizer
-        """Object implementing the Tokenizer protocol."""
-
-    def __call__(
-        self,
-        dataset: T,
-        seq_feature: str,
-        cond_features: list[str],
-        padding: bool,
-        add_special_tokens: bool,
-        device: torch.device | None,
-        **kwargs,
-    ) -> Collated:
-        return self.collate_data(
-            dataset,
-            seq_feature,
-            cond_features,
-            padding,
-            add_special_tokens,
-            device,
-            **kwargs,
-        )
-
-    @abstractmethod
-    def collate_data(
-        self,
-        dataset: T,
-        seq_feature: str,
-        cond_features: list[str],
-        padding: bool,
-        add_special_tokens: bool,
-        device: torch.device | None,
-        **kwargs,
-    ) -> Collated:
-        """Interface method for implementing data collation.
-
-        Args:
-            dataset: Dataset object of target data type.
-            seq_feature: Name of sequence feature or column.
-            cond_features: List of names of condition features or columns.
-            padding: Whether or not to uniformly pad outputs to the maximum length specified in the model configuration.
-            add_special_tokens: Whether or not to add special tokens according to a template. Special tokens will be
-                added by default during training, but will be omitted during inference, where the specified bos_token_id
-                will be appended before each pass through the decoder.
-            device: Torch device to which the produced tensors are moved.
-            **kwargs: Keyword arguments.
-
-        Returns:
-            Collated: Return schema representing a store of tensors ready for direct input to a Model.
-        """
-        ...
-
-
-class DataframeCollator(Collator):
-    """Collator designed for pandas.DataFrame inputs."""
-
-    def collate_data(
-        self,
-        dataset: DataFrame,
-        seq_feature: str,
-        cond_features: list[str],
-        padding: bool,
-        add_special_tokens: bool,
-        device: torch.device | None,
-        **kwargs,
-    ) -> Collated:
-        """Implementation of data collation for pandas.DataFrame inputs.
-
-        Args:
-            dataset: Dataset object of pandas.DataFrame type.
-            seq_feature: Name of sequence column.
-            cond_features: List of names of condition columns.
-            padding: Whether or not to uniformly pad outputs to the maximum length specified in the model configuration.
-            add_special_tokens: Whether or not to add special tokens according to a template. Special tokens will be
-                added by default during training, but will be omitted during inference, where the specified bos_token_id
-                will be appended before each pass through the decoder.
-            device: Torch device to which the produced tensors are moved.
-            **kwargs: Keyword arguments passed to Tokenizer.encode.
-
-        Returns:
-            Collated: Return schema representing a store of tensors ready for direct input to a Model.
-        """
-        df = dataset.copy()
-        input_ids = self.tokenizer.encode(
-            seq=df[seq_feature],
-            padding=padding,
-            max_length=self.model_config.input.max_len,
-            add_special_tokens=add_special_tokens,
-            **kwargs,
-        )
-        conditions = torch.from_numpy(df[cond_features].to_numpy()).to(
-            device=device, dtype=torch.float
-        )
-        token_mask = (input_ids != self.model_config.common.padding_idx).to(torch.bool)
-        condition_mask = (conditions != self.model_config.common.padding_value).to(
-            torch.bool
-        )
-        return Collated(
-            input_ids=input_ids,
-            conditions=conditions,
-            token_mask=token_mask,
-            condition_mask=condition_mask,
-        )

kernel_elastic_autoencoder-2.0.0/src/kernel_elastic_autoencoder/pipeline.py

@@ -1,203 +0,0 @@
-from collections.abc import Callable, Iterable
-
-import torch
-from pydantic import BaseModel, ConfigDict
-
-from kernel_elastic_autoencoder.collate import Collator, DataframeCollator
-from kernel_elastic_autoencoder.config import ModelConfig
-from kernel_elastic_autoencoder.model import Model
-from kernel_elastic_autoencoder.sample import Sampler, Top1Sampler
-from kernel_elastic_autoencoder.tokenizer import Tokenizer
-
-
-class CompletionIntermediate(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    latents: torch.Tensor
-    input_ids: torch.Tensor
-    batches_completed: torch.Tensor
-    condition_embeddings: torch.Tensor
-    condition_mask: torch.Tensor
-
-
-class Completion(BaseModel):
-    """Return schema for Pipeline.completion."""
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    outputs: Iterable[str]
-    """Iterable of generated output sequences."""
-    condition_embeddings: torch.Tensor
-    """Tensor containing generated condition embeddings for reuse."""
-    condition_mask: torch.Tensor
-    """Tensor containing generated condition masks for reuse."""
-
-
-class Pipeline[T]:
-    """User-facing pipeline for inference.
-
-    Defines an easy-to-use API for decoder-only inference with a pretrained model.
-
-    Examples:
-        Wrapping a pretrained model and tokenizer:
-        >>> model = Model.from_pretrained("./checkpoint")
-        >>> tokenizer = MyClassImplementingTokenizer.from_pretrained("./checkpoint/tokenizer")
-        >>> pipe = Pipeline(model, tokenizer, MyClassImplementingCollator, MyClassImplementingSampler)
-        >>> compl = pipe.completion(latents, dataset, seq_feature, cond_features)
-        >>> print(compl.outputs)
-
-        TODO: Complete examples.
-    """
-
-    def __init__(
-        self,
-        model: Model,
-        tokenizer: Tokenizer,
-        collator: Callable[[ModelConfig, Tokenizer], Collator[T]] = DataframeCollator,
-        sampler: Callable[[Tokenizer], Sampler] = Top1Sampler,
-        device: torch.device | None = None,
-    ) -> None:
-        """Instantiates a Pipeline object.
-
-        TODO: Pipeline setup will be streamlined such that collators and samplers are by default fetched from model
-            configuration.
-
-        Args:
-            model: Pre-trained Model object for inference. Can be obtained from Model.from_pretrained.
-            tokenizer: Pre-configured Tokenizer object. The Tokenizer protocol supports tokenizers inheriting from
-                transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded from HuggingFace Hub.
-            collator: Constructor for a data collator inheriting from Collator. Determines input types to Pipeline methods.
-            sampler: Constructor for a data sampler inheriting from Sampler. Determines how outputs are produced from
-                Pipeline methods.
-            device: Torch device used for inference.
-        """
-        self.model = model.to(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. The Tokenizer Protocol supports tokenizers inheriting from
-            transformers.PreTrainedTokenizerBase, so such tokenizers may be loaded from HuggingFace Hub."""
-        self.device = device
-        """Torch device used for inference."""
-
-        self.collator = collator(model.config_typed, tokenizer)
-        """Data collator inheriting from Collator, instantiated by passing Pipeline.model.config_typed and Pipeline.tokenizer."""
-        self.sampler = sampler(tokenizer)
-        """Data sampler inheriting from Sampler, instantiated by passing Pipeline.tokenizer."""
-
-    @torch.inference_mode()
-    def _completion_entry(
-        self,
-        latents: torch.Tensor,
-        dataset: T,
-        seq_feature: str,
-        cond_features: list[str],
-        device: torch.device | None = None,
-        **kwargs,
-    ) -> CompletionIntermediate:
-        ds = self.collator(
-            dataset=dataset,
-            seq_feature=seq_feature,
-            cond_features=cond_features,
-            padding=False,
-            add_special_tokens=False,
-            device=device,
-            **kwargs,
-        )
-        conds_embed = self.model.embed_conditions(ds.conditions)
-        return CompletionIntermediate(
-            latents=latents,
-            input_ids=ds.input_ids,
-            batches_completed=torch.zeros(ds.input_ids.size(0), dtype=torch.bool),
-            condition_embeddings=conds_embed,
-            condition_mask=ds.condition_mask,
-        )
-
-    @torch.inference_mode()
-    def _completion_step(
-        self, intermediate: CompletionIntermediate, **kwargs
-    ) -> CompletionIntermediate:
-        intermediate.input_ids = torch.cat(
-            [
-                torch.full(
-                    (intermediate.input_ids.size(0), 1), self.tokenizer.bos_token_id
-                ),
-                intermediate.input_ids,
-            ],
-            dim=1,
-        )
-        logits = self.model.decode(
-            current_output=intermediate.input_ids,
-            latents=intermediate.latents,
-            condition_embeddings=intermediate.condition_embeddings,
-            token_mask=None,
-            condition_mask=intermediate.condition_mask,
-        )
-        new_toks = self.sampler.sample_ids(logits[:, -1:])
-        intermediate.batches_completed |= (
-            new_toks.squeeze(-1) == self.tokenizer.eos_token_id
-        )
-        torch.where(
-            intermediate.batches_completed.unsqueeze(-1),
-            self.tokenizer.pad_token_id,
-            new_toks,
-        )
-        intermediate.input_ids = torch.cat([intermediate.input_ids, new_toks], dim=1)
-        return intermediate
-
-    @torch.inference_mode()
-    def _completion_exit(
-        self, intermediate: CompletionIntermediate, **kwargs
-    ) -> Completion:
-        return Completion(
-            outputs=self.sampler(
-                intermediate.input_ids, skip_special_tokens=True
-            ),
-            condition_embeddings=intermediate.condition_embeddings,
-            condition_mask=intermediate.condition_mask,
-        )
-
-    def completion(
-        self,
-        latents: torch.Tensor,
-        dataset: T,
-        seq_feature: str,
-        cond_features: list[str],
-        device: torch.device | None = None,
-        **kwargs,
-    ) -> Completion:
-        """Completes each batch in an input dataset with latent and condition guidance through the decoder.
-
-        Using decoder-only inference, the model completes each batch in the provided dataset. Dataset input typing is
-        uniquely defined by the used Collator. The input dataset is passed through the collator, then subject to
-        autoregressive decoder-only inference through the model, where new tokens are sampled from intermediates and
-        stored in a tensor format. Finally, on reaching end-of-sequence tokens or the maximum sequence length, the
-        sequence is decoded by the tokenizer and returned along with the generated condition embeddings and masks in a
-        Completion schema.
-
-        TODO: Callbacks are planned to enable easy custom inference functions and streaming of text.
-
-        Args:
-            latents: Tensor of dimension (B, P * E) containing latent vectors for the batches.
-            dataset: Dataset of type specified in Pipeline.collator.
-            seq_feature: String to pass to collator containing the feature or column containing text sequences to
-                complete.
-            cond_features: List of strings to pass to collator containing features or columns containing numerical
-                condition values.
-            device: Torch device used for inference.
-            **kwargs: Additional keyword arguments passed to Pipeline.collator.
-
-        Returns:
-            Completion: Completion object, containing outputs, as well as condition embeddings and masks for reuse.
-        """
-        intermediate = self._completion_entry(
-            latents=latents,
-            dataset=dataset,
-            seq_feature=seq_feature,
-            cond_features=cond_features,
-            device=device,
-            **kwargs,
-        )
-        while (
-            intermediate.input_ids.size(1) < self.model.config_typed.input.max_len
-        ) and (not intermediate.batches_completed.all()):
-            intermediate = self._completion_step(intermediate=intermediate)
-        return self._completion_exit(intermediate=intermediate)

kernel_elastic_autoencoder-2.0.0/src/kernel_elastic_autoencoder/sample.py

@@ -1,70 +0,0 @@
-from abc import ABC, abstractmethod
-from collections.abc import Iterable
-
-import torch
-
-from kernel_elastic_autoencoder.tokenizer import Tokenizer
-
-
-class Sampler(ABC):
-    """Base class for samplers.
-
-    Provides a specification for sampling text sequences from logits.
-    """
-
-    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]:
-        """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 in intermediate steps.
-
-        Args:
-            logits: Tensor of dimension (B, S, L) containing logits.
-            **kwargs: Keyword arguments.
-
-        Returns:
-            torch.Tensor: Tensor of dimension (B, S) containing vocabulary indices.
-        """
-        ...
-
-
-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.
-
-        Args:
-            logits: Tensor of dimension (B, S, L) containing logits.
-            **kwargs: Keyword arguments.
-
-        Returns:
-            torch.Tensor: Tensor of dimension (B, S) containing vocabulary indices.
-        """
-        return (
-            torch.topk(logits, k=1, dim=-1, **kwargs).indices.squeeze(-1).to(torch.long)
-        )

kernel_elastic_autoencoder-2.0.0/src/kernel_elastic_autoencoder/training.py

@@ -1,386 +0,0 @@
-import os
-from typing import Protocol
-
-import torch
-import torch.distributed as dist
-from pydantic import BaseModel, ConfigDict
-from torch import nn
-from torch.nn.parallel import DistributedDataParallel as DDP
-from torch.utils.data import DistributedSampler
-
-from kernel_elastic_autoencoder.collate import Collated
-from kernel_elastic_autoencoder.config import TrainingConfig
-from kernel_elastic_autoencoder.losses import Loss
-from kernel_elastic_autoencoder.model import Model
-
-
-def _is_local():
-    return (
-        dist.is_torchelastic_launched() and (dist.get_rank() == 0)
-    ) or not dist.is_torchelastic_launched()
-
-
-class TrainerCallbackCtx(BaseModel):
-    model_config = ConfigDict(arbitrary_types_allowed=True)
-    dist: bool
-    device_type: str
-    local_rank: int | None
-    model: Model | DDP
-    optimizer: torch.optim.Optimizer
-    scheduler: torch.optim.lr_scheduler.LRScheduler
-    dataloader_train: torch.utils.data.DataLoader
-    dataloader_test: torch.utils.data.DataLoader
-
-    epoch: int | None = None
-    rel_epoch: int | None = None
-    batch: int | None = None
-    train_loss: torch.Tensor | None = None
-    test_loss: torch.Tensor | None = None
-
-
-class TrainerCallback(Protocol):
-    """Protocol to be implemented for callback classes to a Trainer.
-
-    On each hook, except before the Trainer is initialized, a TrainerCallbackCtx schema is passed with the appropriate
-    information passed.
-    """
-
-    def before_init(self):
-        """Called before training setup."""
-        ...
-
-    def after_init(self, ctx: TrainerCallbackCtx):
-        """Called after training setup."""
-        ...
-
-    def before_epoch(self, ctx: TrainerCallbackCtx):
-        """Called before each epoch."""
-        ...
-
-    def before_train_batch(self, ctx: TrainerCallbackCtx):
-        """Called before each forward pass of a single training batch."""
-        ...
-
-    def after_train_batch(self, ctx: TrainerCallbackCtx):
-        """Called after each forward pass of a single training batch."""
-        ...
-
-    def before_test_batch(self, ctx: TrainerCallbackCtx):
-        """Called before each forward pass of a single test batch."""
-        ...
-
-    def after_test_batch(self, ctx: TrainerCallbackCtx):
-        """Called after each forward pass of a single test batch."""
-        ...
-
-    def after_epoch(self, ctx: TrainerCallbackCtx):
-        """Called after each epoch."""
-        ...
-
-    def after_training(self, ctx: TrainerCallbackCtx):
-        """Called after training ends."""
-        ...
-
-
-class TrainerDefaultCallback:
-    # TODO: Implement sensible default logging.
-    def before_init(self):
-        pass
-
-    def after_init(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def before_epoch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def before_train_batch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def after_train_batch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def before_test_batch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def after_test_batch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def after_epoch(self, ctx: TrainerCallbackCtx):
-        pass
-
-    def after_training(self, ctx: TrainerCallbackCtx):
-        pass
-
-
-class Trainer:
-    def __init__(
-        self,
-        config: dict | TrainingConfig,
-        callbacks: tuple[TrainerCallback] = (TrainerDefaultCallback(),),
-    ):
-        """Instantiates a Trainer object.
-
-        Args:
-            config: Dictionary or TrainingConfig schema defining model parameters. Will be validated with TrainingConfig
-                regardless of input type.
-            callbacks: Tuple of callback classes implementing TrainerCallback.
-        """
-        self.config = config
-        """Configuration object for Hugging Face Hub compatible serialization. Not recommended to use, as it is 
-                internal-use. Use Trainer.config_typed instead."""
-        config_typed = TrainingConfig.model_validate(config)
-        self.config_typed: TrainingConfig = config_typed
-        """Type-validated config in a Pydantic TrainingConfig schema, recommended for public API use."""
-        self.callbacks = callbacks
-        """Tuple of callback classes implementing TrainerCallback."""
-
-        self._ctx: TrainerCallbackCtx
-
-        self._model: DDP | Model
-        self._dist: bool
-        self._loss_fn: nn.Module
-        self._optimizer: torch.optim.Optimizer
-        self._scheduler: torch.optim.lr_scheduler.LRScheduler
-        self._dataloader_train: torch.utils.data.DataLoader
-        self._dataloader_test: torch.utils.data.DataLoader
-        self._device_type: str
-        self._local_rank: int | None
-
-    def train(
-        self,
-        model: Model,
-        ds: Collated,
-        train_split: float = 0.9,
-        checkpoint: str = "./checkpoint",
-    ):
-        """Trains a model with a Collated dataset. Optionally resumes from an existing checkpoint.
-
-        Args:
-            model: Freshly instantiated model.
-            ds: Tensor dataset following the Collated schema.
-            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.
-        """
-        next_epoch = self._setup(model, ds, train_split, checkpoint)
-        for epoch in range(next_epoch, self.config_typed.common.max_epochs):
-            self._ctx.rel_epoch = epoch - next_epoch
-            self._epoch(epoch, checkpoint)
-        self._ctx.epoch = None
-        self._ctx.rel_epoch = None
-        if _is_local():
-            for cb in self.callbacks:
-                cb.after_training(self._ctx)
-
-    def _setup(
-        self,
-        model: Model,
-        ds: Collated,
-        train_split: float,
-        checkpoint: str,
-    ):
-        if _is_local():
-            for cb in self.callbacks:
-                cb.before_init()
-        self._model = model
-        self._dist = dist.is_torchelastic_launched()
-        if self._dist:
-            self._model, self._device_type, self._local_rank = self._setup_ddp()
-        else:
-            self._device_type, self._local_rank = self._setup_no_ddp()
-        self._loss_fn = self._setup_loss(model)
-        self._optimizer, self._scheduler = self._setup_optimizer()
-        self._dataloader_train, self._dataloader_test = self._setup_dataloaders(
-            ds, train_split
-        )
-        next_epoch = 0
-        if os.path.exists(checkpoint):
-            next_epoch = self._resume(checkpoint)
-        self._ctx = TrainerCallbackCtx(
-            dist=self._dist,
-            device_type=self._device_type,
-            local_rank=self._local_rank,
-            model=self._model,
-            optimizer=self._optimizer,
-            scheduler=self._scheduler,
-            dataloader_train=self._dataloader_train,
-            dataloader_test=self._dataloader_test,
-        )
-        if _is_local():
-            for cb in self.callbacks:
-                cb.after_init(self._ctx)
-        return next_epoch
-
-    def _setup_loss(self, model: Model):
-        _loss_fn = Loss(
-            hp_lambda=self.config_typed.hyperparameters.hp_lambda,
-            hp_delta=self.config_typed.hyperparameters.hp_delta,
-            hp_sigma=self.config_typed.hyperparameters.hp_sigma,
-            kernel_dist_size=self.config_typed.hyperparameters.kernel_dist_size,
-            padding_idx=model.config_typed.common.padding_idx,
-            embedding_dim=model.config_typed.common.embedding_dim,
-            pooling_dim=model.config_typed.common.pooling_dim,
-        )
-        return _loss_fn
-
-    def _setup_optimizer(self):
-        optimizer = self.config_typed.optimizer.optimizer_fn(
-            self._model.parameters(), **self.config_typed.optimizer.optimizer_params
-        )
-        scheduler = self.config_typed.optimizer.scheduler_fn(
-            optimizer, **self.config_typed.optimizer.scheduler_params
-        )
-        return optimizer, scheduler
-
-    def _setup_ddp(self) -> tuple[DDP, str, int]:
-        device_type, vendor_backend = self._get_backend()
-        dist.init_process_group(backend=vendor_backend)
-        local_rank = int(os.environ["LOCAL_RANK"])
-        model = DDP(self._model.to(local_rank))
-        return model, device_type, local_rank
-
-    def _setup_no_ddp(self):
-        device_type, _ = self._get_backend()
-        return device_type, None
-
-    def _setup_dataloaders(self, ds: Collated, train_split: float):
-        dataset = torch.utils.data.TensorDataset(
-            ds.input_ids, ds.conditions, ds.token_mask, ds.condition_mask
-        )
-        dataset_train, dataset_test = torch.utils.data.random_split(
-            dataset, [train_split, 1 - train_split]
-        )
-        if self._dist:
-            sampler_train = DistributedSampler(dataset_train)
-            sampler_test = DistributedSampler(dataset_test)
-        else:
-            sampler_train = torch.utils.data.RandomSampler(dataset_train)
-            sampler_test = torch.utils.data.SequentialSampler(dataset_test)
-        _dataloader_train = torch.utils.data.DataLoader(
-            dataset_train,
-            batch_size=self.config_typed.common.batch_size,
-            sampler=sampler_train,
-            pin_memory=True,
-        )
-        _dataloader_test = torch.utils.data.DataLoader(
-            dataset_test,
-            batch_size=self.config_typed.common.batch_size,
-            sampler=sampler_test,
-            pin_memory=True,
-        )
-        return _dataloader_train, _dataloader_test
-
-    def _resume(
-        self,
-        checkpoint: str,
-    ):
-        train_state = torch.load(os.path.join(checkpoint, "train_state.pt"))
-        self._optimizer.load_state_dict(train_state["optimizer"])
-        self._scheduler.load_state_dict(train_state["scheduler"])
-        self._model.from_pretrained(checkpoint)  # type: ignore
-        return train_state["next_epoch"]
-
-    def _epoch(
-        self,
-        epoch: int,
-        checkpoint: str,
-    ):
-        self._ctx.epoch = epoch
-        if _is_local():
-            for cb in self.callbacks:
-                cb.before_epoch(self._ctx)
-
-        self._model.train()
-        self._ctx.batch = 0
-        for input_ids, conditions, token_mask, condition_mask in self._dataloader_train:
-            if _is_local():
-                for cb in self.callbacks:
-                    cb.before_train_batch(self._ctx)
-            self._ctx.train_loss = self._batch_train(
-                input_ids, conditions, token_mask, condition_mask
-            )
-            if _is_local():
-                for cb in self.callbacks:
-                    cb.after_train_batch(self._ctx)
-            self._ctx.batch += 1
-
-        self._model.eval()
-        self._ctx.batch = 0
-        for input_ids, conditions, token_mask, condition_mask in self._dataloader_test:
-            if _is_local():
-                for cb in self.callbacks:
-                    cb.before_test_batch(self._ctx)
-            self._ctx.test_loss = self._batch_test(
-                input_ids, conditions, token_mask, condition_mask
-            )
-            if _is_local():
-                for cb in self.callbacks:
-                    cb.after_test_batch(self._ctx)
-            self._ctx.batch += 1
-
-        self._end_epoch(epoch, checkpoint)
-        self._ctx.train_loss = None
-        self._ctx.test_loss = None
-        self._ctx.batch = None
-        if _is_local():
-            for cb in self.callbacks:
-                cb.after_epoch(self._ctx)
-
-    def _end_epoch(self, epoch: int, checkpoint: str):
-        self._scheduler.step(epoch)
-        self._save_state(checkpoint, epoch)
-
-    def _batch_train(self, input_ids, conditions, token_mask, condition_mask):
-        self._optimizer.zero_grad()
-        with torch.amp.autocast(self._device_type):
-            prediction, prediction_noise, latents_noise = self._model.forward(
-                input_ids, conditions, token_mask, condition_mask
-            )
-            loss = self._loss_fn(
-                prediction, prediction_noise, input_ids[:, 1:], latents_noise
-            )
-        loss.backward()
-        self._optimizer.step()
-        return loss.detach()
-
-    def _batch_test(self, input_ids, conditions, token_mask, condition_mask):
-        with torch.no_grad():
-            with torch.amp.autocast(self._device_type):
-                prediction, prediction_noise, latents_noise = self._model.forward(
-                    input_ids,
-                    conditions,
-                    token_mask,
-                    condition_mask,
-                )
-            loss = self._loss_fn(
-                prediction, prediction_noise, input_ids[:, 1:], latents_noise
-            )
-        return loss.detach()
-
-    def _save_state(self, checkpoint: str, epoch: int):
-        if _is_local():
-            os.makedirs(checkpoint, exist_ok=True)
-            self._model.save_pretrained(checkpoint)  # type: ignore
-            torch.save(
-                {
-                    "next_epoch": epoch + 1,
-                    "optimizer": self._optimizer.state_dict(),
-                    "scheduler": self._scheduler.state_dict(),
-                },
-                os.path.join(checkpoint, "train_state.pt"),
-            )
-            self.config_typed.to_json(os.path.join(checkpoint, "train_config.json"))
-
-    def _get_backend(self) -> tuple[str, str]:
-        if torch.accelerator.is_available():
-            device_type = torch.accelerator.current_accelerator().type  # type: ignore
-            vendor_backend = torch.distributed.get_default_backend_for_device(
-                device_type
-            )
-
-        else:
-            device_type = torch.device("cpu").type
-            vendor_backend = torch.distributed.get_default_backend_for_device(
-                device_type
-            )
-
-        return device_type, vendor_backend