deepchopper 1.3.0__cp310-abi3-macosx_11_0_arm64.whl

deepchopper/models/basic_module.py

@@ -0,0 +1,243 @@
+from typing import Any
+
+import torch
+from huggingface_hub import PyTorchModelHubMixin
+from lightning import LightningModule
+from torch import nn
+from torchmetrics import MaxMetric, MeanMetric
+from torchmetrics.classification import F1Score, Precision, Recall
+
+
+class ContinuousIntervalLoss(nn.Module):
+    """A custom loss function that penalizes the model for predicting different classes in consecutive positions."""
+
+    def __init__(self, lambda_penalty: float = 0, **kwargs):
+        super().__init__()
+        self.base = torch.nn.CrossEntropyLoss(**kwargs)
+        self.lambda_penalty = lambda_penalty
+
+    @property
+    def ignore_index(self):
+        return self.base.ignore_index
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        loss = self.base(pred, target)
+        if self.lambda_penalty == 0:
+            return loss
+        valid_mask = target != self.ignore_index
+        true_pred = pred.argmax(-1)[valid_mask]
+        true_target = target[valid_mask]
+        penalty = self.lambda_penalty * (true_pred[1:] != true_target[:-1]).float().mean()
+        return loss + penalty
+
+
+class TokenClassificationLit(LightningModule, PyTorchModelHubMixin):
+    """A PyTorch Lightning module for training a token classification model."""
+
+    def __init__(
+        self,
+        net: nn.Module,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler.LRScheduler,
+        criterion: nn.Module,
+        *,
+        compile: bool,
+    ):
+        """Genomics Benchmark CNN model for PyTorch Lightning.
+
+        :param net: The CNN model.
+        :param scheduler: The learning rate scheduler to use for training.
+        """
+        super().__init__()
+
+        self.example_input_array = {
+            "input_ids": torch.randint(0, 11, (1, 1000)),
+            "input_quals": torch.rand(1, 1000),
+        }  # [batch, seq_len]
+
+        # this line allows to access init params with 'self.hparams' attribute
+        # also ensures init params will be stored in ckpt
+        self.save_hyperparameters(logger=False, ignore=["net", "criterion"])
+        self.net = net
+        # loss function
+        self.criterion = criterion
+
+        # metric objects for calculating and averaging accuracy across batches
+        self.train_acc = F1Score(
+            task="binary", num_classes=net.number_of_classes, ignore_index=self.criterion.ignore_index
+        )
+        self.val_acc = F1Score(
+            task="binary", num_classes=net.number_of_classes, ignore_index=self.criterion.ignore_index
+        )
+        self.test_acc = F1Score(
+            task="binary", num_classes=net.number_of_classes, ignore_index=self.criterion.ignore_index
+        )
+
+        self.test_precision = Precision(
+            task="binary", num_classes=net.number_of_classes, ignore_index=self.criterion.ignore_index
+        )
+        self.test_recall = Recall(
+            task="binary", num_classes=net.number_of_classes, ignore_index=self.criterion.ignore_index
+        )
+
+        # for averaging loss across batches
+        self.train_loss = MeanMetric()
+        self.val_loss = MeanMetric()
+        self.test_loss = MeanMetric()
+        # for tracking best so far validation accuracy
+        self.val_acc_best = MaxMetric()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        input_quals: torch.Tensor,
+    ) -> torch.Tensor:
+        """Perform a forward pass through the model `self.net`.
+
+        :param x: A tensor of images.
+        :return: A tensor of logits.
+        """
+        return self.net(input_ids, input_quals)
+
+    def on_train_start(self) -> None:
+        """Lightning hook that is called when training begins."""
+        # by default lightning executes validation step sanity checks before training starts,
+        # so it's worth to make sure validation metrics don't store results from these checks
+        self.val_loss.reset()
+        self.val_acc.reset()
+        self.val_acc_best.reset()
+
+    def model_step(self, batch: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Perform a single model step on a batch of data.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target labels.
+        :return: A tuple containing (in order):
+            - A tensor of losses.
+            - A tensor of predictions.
+            - A tensor of target labels.
+        """
+        input_ids = batch["input_ids"]
+        input_quals = batch["input_quals"]
+        logits = self.forward(input_ids, input_quals)
+        loss = self.criterion(logits.reshape(-1, logits.size(-1)), batch["labels"].long().view(-1))
+        preds = torch.argmax(logits, dim=-1)
+        return loss, preds, batch["labels"]
+
+    def training_step(self, batch: dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Perform a single training step on a batch of data from the training set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        :return: A tensor of losses between model predictions and targets.
+        """
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.train_loss(loss)
+        self.train_acc(preds, targets)
+
+        self.log("train/loss", self.train_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("train/f1", self.train_acc, on_step=False, on_epoch=True, prog_bar=True)
+
+        # return loss or backpropagation will fail
+        return loss
+
+    def on_train_epoch_end(self) -> None:
+        """Lightning hook that is called when a training epoch ends."""
+
+    def validation_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
+        """Perform a single validation step on a batch of data from the validation set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        """
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.val_loss(loss)
+        self.val_acc(preds, targets)
+
+        self.log("val/loss", self.val_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("val/f1", self.val_acc, on_step=False, on_epoch=True, prog_bar=True)
+
+    def on_validation_epoch_end(self) -> None:
+        """Lightning hook that is called when a validation epoch ends."""
+        acc = self.val_acc.compute()  # get current val acc
+        self.val_acc_best(acc)  # update best so far val acc
+        # log `val_acc_best` as a value through `.compute()` method, instead of as a metric object
+        # otherwise metric would be reset by lightning after each epoch
+        self.log("val/f1_best", self.val_acc_best.compute(), sync_dist=True, prog_bar=True)
+
+    def test_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> None:
+        """Perform a single test step on a batch of data from the test set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        """
+        loss, preds, targets = self.model_step(batch)
+
+        # update and log metrics
+        self.test_loss(loss)
+        self.test_acc(preds, targets)
+
+        self.test_precision(preds, targets)
+        self.test_recall(preds, targets)
+
+        self.log("test/loss", self.test_loss, on_step=False, on_epoch=True, prog_bar=True)
+        self.log("test/f1", self.test_acc, on_step=False, on_epoch=True, prog_bar=True)
+
+    def on_test_epoch_end(self) -> None:
+        """Lightning hook that is called when a test epoch ends."""
+        self.log("test/precision", self.test_precision)
+        self.log("test/recall", self.test_recall)
+
+    def predict_step(self, batch: tuple[torch.Tensor, torch.Tensor], batch_idx: int) -> torch.Tensor:
+        """Perform a single prediction step on a batch of data from the test set.
+
+        :param batch: A batch of data (a tuple) containing the input tensor of images and target
+            labels.
+        :param batch_idx: The index of the current batch.
+        """
+        input_ids = batch["input_ids"]
+        input_quals = batch["input_quals"]
+        logits = self.forward(input_ids, input_quals)
+        return logits, batch["labels"]
+
+    def setup(self, stage: str) -> None:
+        """Lightning hook that is called at the beginning of fit (train + validate), validate, test, or predict.
+
+        This is a good hook when you need to build models dynamically or adjust something about
+        them. This hook is called on every process when using DDP.
+
+        :param stage: Either `"fit"`, `"validate"`, `"test"`, or `"predict"`.
+        """
+        if self.hparams.compile and stage == "fit":
+            self.net = torch.compile(self.net)
+
+    def configure_optimizers(self) -> dict[str, Any]:
+        """Choose what optimizers and learning-rate schedulers to use in your optimization.
+
+        Normally you'd need one. But in the case of GANs or similar you might have multiple.
+
+        Examples:
+            https://lightning.ai/docs/pytorch/latest/common/lightning_module.html#configure-optimizers
+
+        :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
+        """
+        optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
+
+        if self.hparams.scheduler is not None:
+            scheduler = self.hparams.scheduler(optimizer=optimizer)
+            return {
+                "optimizer": optimizer,
+                "lr_scheduler": {
+                    "scheduler": scheduler,
+                    "monitor": "val/loss",
+                    "interval": "epoch",
+                    "frequency": 1,
+                },
+            }
+        return {"optimizer": optimizer}

deepchopper/models/callbacks.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+
+import torch
+from lightning.pytorch.callbacks import BasePredictionWriter
+
+
+class CustomWriter(BasePredictionWriter):
+    def __init__(self, output_dir, write_interval="epoch"):
+        super().__init__(write_interval)
+        self.output_dir = Path(output_dir)
+
+    def write_on_batch_end(self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx):
+        folder = self.output_dir / str(dataloader_idx)
+        if not folder.exists():
+            folder.mkdir(parents=True, exist_ok=True)
+
+        save_prediction = {
+            "prediction": prediction[0].cpu(),
+            "target": prediction[1].to(torch.int64).cpu(),
+            "seq": batch["input_ids"].cpu(),
+            "qual": batch["input_quals"].cpu(),
+            "id": batch["id"].to(torch.int64).cpu(),
+        }
+
+        torch.save(save_prediction, folder / f"{trainer.global_rank}_{batch_idx}.pt")
+
+    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
+        # WARN: This is a simple implementation that saves all predictions in a single file
+        if not self.output_dir.exists():
+            self.output_dir.mkdir(parents=False, exist_ok=True)
+
+        torch.save(predictions, self.output_dir / "predictions.pt")
+
+
+class PredictionWriter(BasePredictionWriter):
+    def __init__(self, output_dir, write_interval="epoch"):
+        super().__init__(write_interval)
+        self.output_dir = Path(output_dir)
+
+    def write_on_batch_end(self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx):
+        folder = self.output_dir / str(dataloader_idx)
+        if not folder.exists():
+            folder.mkdir(parents=True, exist_ok=True)
+
+        save_prediction = {
+            "prediction": prediction[0].cpu(),
+            "id": batch["id"].to(torch.int64).cpu(),
+        }
+
+        torch.save(save_prediction, folder / f"{trainer.global_rank}_{batch_idx}.pt")
+
+    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
+        # WARN: This is a simple implementation that saves all predictions in a single file
+        if not self.output_dir.exists():
+            self.output_dir.mkdir(parents=False, exist_ok=True)
+
+        torch.save(predictions, self.output_dir / "predictions.pt")

deepchopper/models/cnn.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+
+
+class BenchmarkCNN(nn.Module):
+    """BenchmarkCNN."""
+
+    def __init__(self, number_of_classes, vocab_size, num_filters, filter_sizes, embedding_dim=100):
+        """Genomics Benchmark CNN model.
+
+        `embedding_dim` = 100 comes from:
+        https://github.com/ML-Bioinfo-CEITEC/genomic_benchmarks/tree/main/experiments/torch_cnn_experiments
+        """
+        super().__init__()
+        self.number_of_classes = number_of_classes
+        self.embeddings = nn.Embedding(vocab_size, embedding_dim)
+
+        self.qual_linear1 = nn.Sequential(
+            nn.Linear(1, embedding_dim),
+        )
+
+        layers = []
+        in_channels = embedding_dim
+        for idx, fs in enumerate(filter_sizes):
+            layers.append(
+                nn.Conv1d(
+                    in_channels=in_channels,
+                    out_channels=num_filters[idx],
+                    kernel_size=fs,
+                    padding="same",
+                )
+            )
+            layers.append(nn.BatchNorm1d(num_filters[idx]))
+            layers.append(nn.ReLU())
+            in_channels = num_filters[idx]
+
+        self.model = nn.Sequential(*layers)
+
+        # use number of kernel same as the length of the sequence and average pooling
+        # then flatten and use dense layers
+        self.dense_model = nn.Sequential(
+            nn.Linear(in_channels, number_of_classes),
+        )
+
+    def forward(
+        self, input_ids: torch.Tensor, input_quals: torch.Tensor
+    ):  # Adding `state` to be consistent with other models
+        x = self.embeddings(input_ids)
+        x = F.relu(x + self.qual_linear1(input_quals.unsqueeze(-1)))
+        x = x.transpose(1, 2)  # [batch_size, embedding_dim, input_len]
+        x = self.model(x)
+        x = x.transpose(1, 2)
+        return self.dense_model(x)  # [batch_size, input_len, num_filters]

deepchopper/models/components/__init__.py

@@ -0,0 +1 @@
+"""model components."""

deepchopper/models/dc_hg.py

@@ -0,0 +1,163 @@
+from functools import partial
+
+import torch
+
+from . import basic_module, llm
+from .basic_module import TokenClassificationLit
+
+
+class DeepChopper:
+    """DeepChopper: A genomic language model for chimera artifact detection.
+
+    This class provides convenient methods to load DeepChopper models from checkpoints or
+    from pretrained models on the Hugging Face Hub, and to push trained models to the Hub.
+
+    Example:
+        Load a pretrained model:
+        >>> model = DeepChopper.from_pretrained("yangliz5/deepchopper")
+
+        Load from a local checkpoint:
+        >>> model = DeepChopper.from_checkpoint("path/to/checkpoint.ckpt")
+
+        Push a model to Hugging Face Hub:
+        >>> model = DeepChopper.to_hub("username/model-name", "path/to/checkpoint.ckpt")
+    """
+
+    @staticmethod
+    def to_hub(
+        model_name: str,
+        checkpoint_path: str,
+        *,
+        commit_message: str = "Upload DeepChopper model",
+        private: bool = False,
+        token: str | None = None,
+    ):
+        """Load a model from a checkpoint and push it to the Hugging Face Hub.
+
+        Args:
+            model_name: The repository ID on Hugging Face Hub (format: username/model-name)
+            checkpoint_path: Path to the local checkpoint file (.ckpt)
+            commit_message: Commit message for the upload (default: "Upload DeepChopper model")
+            private: Whether to create a private repository (default: False)
+            token: Hugging Face API token. If None, uses the stored token from `huggingface-cli login`
+
+        Returns:
+            The loaded TokenClassificationLit model
+
+        Example:
+            >>> model = DeepChopper.to_hub(
+            ...     "username/deepchopper-v1",
+            ...     "epoch_012_f1_0.9947.ckpt",
+            ...     commit_message="Upload DeepChopper v1.0",
+            ...     private=False
+            ... )
+        """
+        model = DeepChopper.from_checkpoint(checkpoint_path)
+
+        # Prepare kwargs for push_to_hub
+        push_kwargs = {
+            "repo_id": model_name,
+            "commit_message": commit_message,
+            "private": private,
+        }
+
+        if token is not None:
+            push_kwargs["token"] = token
+
+        model.push_to_hub(**push_kwargs)
+        return model
+
+    @staticmethod
+    def from_checkpoint(checkpoint_path: str):
+        """Load a DeepChopper model from a local checkpoint file.
+
+        This method creates a new TokenClassificationLit model with the HyenaDNA backbone
+        and loads the weights from the specified checkpoint file.
+
+        Args:
+            checkpoint_path: Path to the checkpoint file (.ckpt) containing model weights
+
+        Returns:
+            A TokenClassificationLit model loaded with the checkpoint weights
+
+        Example:
+            >>> model = DeepChopper.from_checkpoint("epoch_012_f1_0.9947.ckpt")
+
+        Note:
+            This function loads checkpoints with weights_only=False to support PyTorch Lightning
+            checkpoints containing optimizer/scheduler configs. Only load checkpoints from trusted sources.
+        """
+        model = TokenClassificationLit(
+            net=llm.hyena.TokenClassificationModule(
+                number_of_classes=2,
+                backbone_name="hyenadna-small-32k-seqlen",
+                freeze_backbone=False,
+                head=llm.TokenClassificationHead(
+                    input_size=256,
+                    lin1_size=1024,
+                    lin2_size=1024,
+                    num_class=2,
+                    use_identity_layer_for_qual=True,
+                    use_qual=True,
+                ),
+            ),
+            optimizer=partial(  # type: ignore[arg-type]
+                torch.optim.Adam,
+                lr=0.0002,
+                weight_decay=0,
+            ),
+            scheduler=partial(torch.optim.lr_scheduler.ReduceLROnPlateau, mode="min", factor=0.1, patience=10),  # type: ignore[arg-type]
+            criterion=basic_module.ContinuousIntervalLoss(lambda_penalty=0),
+            compile=False,
+        )
+        # weights_only=False is required for PyTorch Lightning checkpoints that contain
+        # optimizer/scheduler configs (functools.partial). Only use with trusted checkpoints.
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+        model.load_state_dict(checkpoint["state_dict"])
+        return model
+
+    @staticmethod
+    def from_pretrained(model_name: str):
+        """Load a pretrained DeepChopper model from the Hugging Face Hub.
+
+        This method downloads and loads a pretrained model from the Hugging Face Hub.
+        The model architecture is automatically configured to match the expected
+        HyenaDNA-based token classification setup.
+
+        Args:
+            model_name: The repository ID on Hugging Face Hub (e.g., "yangliz5/deepchopper")
+
+        Returns:
+            A TokenClassificationLit model loaded with pretrained weights
+
+        Example:
+            >>> model = DeepChopper.from_pretrained("yangliz5/deepchopper")
+
+        Note:
+            This requires an internet connection to download the model from Hugging Face Hub.
+            The model will be cached locally after the first download.
+        """
+        return TokenClassificationLit.from_pretrained(
+            model_name,
+            net=llm.hyena.TokenClassificationModule(
+                number_of_classes=2,
+                backbone_name="hyenadna-small-32k-seqlen",
+                freeze_backbone=False,
+                head=llm.TokenClassificationHead(
+                    input_size=256,
+                    lin1_size=1024,
+                    lin2_size=1024,
+                    num_class=2,
+                    use_identity_layer_for_qual=True,
+                    use_qual=True,
+                ),
+            ),
+            optimizer=partial(  # type: ignore[arg-type]
+                torch.optim.Adam,
+                lr=0.0002,
+                weight_decay=0,
+            ),
+            scheduler=partial(torch.optim.lr_scheduler.ReduceLROnPlateau, mode="min", factor=0.1, patience=10),  # type: ignore[arg-type]
+            criterion=basic_module.ContinuousIntervalLoss(lambda_penalty=0),
+            compile=False,
+        )

deepchopper/models/llm/__init__.py

@@ -0,0 +1,32 @@
+"""hyena model components and utilities."""
+
+from .components import (
+    HyenadnaMaxLengths,
+    TokenClassification,
+    TokenClassificationConfig,
+)
+from .head import TokenClassificationHead
+from .hyena import TokenClassificationModule
+from .metric import IGNORE_INDEX, compute_metrics
+from .tokenizer import (
+    DataCollatorForTokenClassificationWithQual,
+    load_tokenizer_from_hyena_model,
+    tokenize_and_align_labels_and_quals,
+    tokenize_and_align_labels_and_quals_ids,
+    tokenize_dataset,
+)
+
+__all__ = [
+    "IGNORE_INDEX",
+    "DataCollatorForTokenClassificationWithQual",
+    "HyenadnaMaxLengths",
+    "TokenClassification",
+    "TokenClassificationConfig",
+    "TokenClassificationHead",
+    "TokenClassificationModule",
+    "compute_metrics",
+    "load_tokenizer_from_hyena_model",
+    "tokenize_and_align_labels_and_quals",
+    "tokenize_and_align_labels_and_quals_ids",
+    "tokenize_dataset",
+]

deepchopper/models/llm/caduceus.py

@@ -0,0 +1,55 @@
+import torch
+from torch import nn
+from transformers import AutoModel
+
+BACKBONES = [
+    "hyenadna-tiny-1k-seqlen",
+    "hyenadna-small-32k-seqlen",
+    "hyenadna-medium-160k-seqlen",
+    "hyenadna-medium-450k-seqlen",
+    "hyenadna-large-1m-seqlen",
+    "caduceus-ph_seqlen-131k_d_model-256_n_layer-16",
+    "caduceus-ps_seqlen-131k_d_model-256_n_layer-16",
+]
+
+# https://github.com/kuleshov-group/caduceus
+
+
+class TokenClassificationModule(nn.Module):
+    """Token classification model."""
+
+    def __init__(
+        self,
+        number_of_classes: int,
+        head: nn.Module,
+        backbone_name: str = "caduceus-ph_seqlen-131k_d_model-256_n_layer-16",
+    ):
+        super().__init__()
+        self.number_of_classes = number_of_classes
+        self.backbone_name = backbone_name
+
+        if "hyenadna" in backbone_name:
+            model_name = f"LongSafari/{backbone_name}-hf"
+        elif "caduceus" in backbone_name:
+            model_name = f"kuleshov-group/{backbone_name}"
+        else:
+            msg = f"Unknown backbone model: {backbone_name}"
+            raise ValueError(msg)
+
+        self.backbone = AutoModel.from_pretrained(model_name, trust_remote_code=True)
+        self.head = head
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        input_quals: torch.Tensor,
+    ):
+        transformer_outputs = self.backbone(
+            input_ids,
+            inputs_embeds=None,
+            output_hidden_states=None,
+            return_dict=None,
+        )
+
+        hidden_states = transformer_outputs[0]
+        return self.head(hidden_states, input_quals)