deepchopper 1.3.0__cp310-abi3-macosx_11_0_arm64.whl

deepchopper/models/llm/components.py

@@ -0,0 +1,99 @@
+import torch
+from torch import nn
+from transformers import AutoModel, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import TokenClassifierOutput
+
+from .head import TokenClassificationHead
+
+HyenadnaMaxLengths = {
+    "hyenadna-tiny-1k-seqlen": 1024,
+    "hyenadna-small-32k-seqlen": 32768,
+    "hyenadna-medium-160k-seqlen": 160000,
+    "hyenadna-medium-450k-seqlen": 450000,  # T4 up to here
+    "hyenadna-large-1m-seqlen": 1_000_000,  # only A100 (paid tier)
+}
+
+
+class TokenClassificationConfig(PretrainedConfig):
+    """Configuration class to store the model's hyperparameters."""
+
+    model_type = "token-classification"
+
+    def __init__(
+        self,
+        input_size: int = 256,
+        lin1_size: int = 1024,
+        lin2_size: int = 1024,
+        num_class: int = 2,
+        *,
+        use_identity_layer_for_qual: bool = True,
+        use_qual: bool = True,
+        **kwargs,
+    ):
+        self.input_size = input_size
+        self.lin1_size = lin1_size
+        self.lin2_size = lin2_size
+        self.num_class = num_class
+        self.use_identity_layer_for_qual = use_identity_layer_for_qual
+        self.use_qual = use_qual
+        super().__init__(**kwargs)
+
+
+class TokenClassification(PreTrainedModel):
+    """Token classification model."""
+
+    config_class = TokenClassificationConfig
+
+    def __init__(
+        self,
+        config,
+        hyenadna_model: str = "hyenadna-small-32k-seqlen",
+        **kwargs,
+    ):
+        super().__init__(config, **kwargs)
+        self.num_class = config.num_class
+        self.hyenadna_model_name = hyenadna_model
+        self.hyenadna = AutoModel.from_pretrained(f"LongSafari/{hyenadna_model}-hf", trust_remote_code=True)
+
+        self.head = TokenClassificationHead(
+            input_size=config.input_size,
+            lin1_size=config.lin1_size,
+            lin2_size=config.lin2_size,
+            num_class=config.num_class,
+            use_identity_layer_for_qual=config.use_identity_layer_for_qual,
+            use_qual=config.use_qual,
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor,
+        input_quals: torch.Tensor,
+        inputs_embeds: torch.FloatTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ):
+        transformer_outputs = self.hyenadna(
+            input_ids,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        _batch_size = input_ids.shape[0]
+        hidden_states = transformer_outputs[0]
+
+        logits = self.head(hidden_states, input_quals)
+        labels = labels.to(logits.device)
+        loss_fct = nn.CrossEntropyLoss()
+
+        loss = loss_fct(logits.view(-1, self.num_class), labels.view(-1))
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=transformer_outputs.hidden_states,
+        )

deepchopper/models/llm/head.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+
+
+class TokenClassificationCnnHead(nn.Module):
+    def __init__(
+        self,
+        input_size,
+        number_of_classes,
+        num_filters,
+        filter_sizes,
+    ):
+        super().__init__()
+        self.number_of_classes = number_of_classes
+
+        self.qual_linear1 = nn.Linear(num_filters, number_of_classes)
+
+        layers = []
+        in_channels = input_size
+
+        for idx, fs in enumerate(filter_sizes):
+            layers.append(nn.Conv1d(in_channels=in_channels, out_channels=num_filters, 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)
+        self.dense = nn.Linear(in_channels, number_of_classes)
+
+    def forward(self, x: torch.Tensor, input_quals: torch.Tensor):
+        x = F.relu(x + self.qual_linear1(input_quals.unsqueeze(-1)))
+        x = x.transpose(1, 2)  # (batch, num_filters, seq_len)
+        x = self.model(x)  # (batch, num_filters, seq_len)
+        x = x.transpose(1, 2)  # (batch, seq_len, num_filters)
+        return self.dense(x)
+
+
+class TokenClassificationHead(nn.Module):
+    """Token classification head for the model."""
+
+    def __init__(
+        self,
+        input_size: int,
+        num_class: int,
+        lin1_size: int,
+        lin2_size: int,
+        *,
+        use_identity_layer_for_qual: bool,
+        use_qual: bool,
+    ):
+        """Initialize the neural network model.
+
+        Parameters:
+            input_size (int): The size of the input features.
+            lin1_size (int): The size of the first linear layer.
+            lin2_size (int): The size of the second linear layer.
+            num_class (int): The number of output classes.
+            use_identity_layer_for_qual (bool): Whether to use an identity layer for quality.
+            use_qual (bool): Whether to use quality in the model.
+        """
+        if lin1_size != lin2_size:
+            msg = f"{lin1_size=} and {lin2_size=} must be equal"
+            raise ValueError(msg)
+
+        super().__init__()
+        self.use_qual = use_qual
+        self.activation = nn.ReLU()
+        self.linear1 = nn.Linear(input_size, lin1_size)
+        self.linear2 = nn.Linear(lin1_size, lin2_size)
+        self.linear3 = nn.Linear(lin2_size, num_class)
+
+        self.qual_linear1 = nn.Identity() if use_identity_layer_for_qual else nn.Linear(1, lin1_size)
+
+    def forward(self, x: torch.Tensor, input_quals: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the neural network model.
+
+        Parameters:
+            x (torch.Tensor): Input tensor to the model.
+            input_quals (torch.Tensor): Input tensor representing qualities.
+
+        Returns:
+            torch.Tensor: Output tensor from the model.
+
+        This method performs a forward pass through the neural network model.
+        It takes in two input tensors, x and input_quals, and processes them through the model layers.
+        The output tensor is returned from the model after passing through the linear and activation layers.
+        If the 'use_qual' flag is set to True, the input_quals tensor is used to calculate a residual value that is added to the output tensor before passing through the linear and activation layers again.
+        This helps incorporate qualities into the model's predictions.
+        If the 'use_qual' flag is set to False, the input_quals tensor is not used and the output tensor from the first linear layer is directly passed through the second linear and activation layers.
+        The final output tensor is returned from the model after passing through the third linear layer.
+        Note: The activation function used in the model is specified by self.activation and should be set during model initialization.
+        """
+        output = self.activation(self.linear1(x))
+
+        if self.use_qual:
+            residual = output + self.qual_linear1(input_quals.unsqueeze(-1))  # may add activation
+            output = self.activation(self.linear2(residual) + residual)
+        else:
+            output = self.activation(self.linear2(output))
+
+        return self.linear3(output)

deepchopper/models/llm/hyena.py

@@ -0,0 +1,41 @@
+import torch
+from torch import nn
+from transformers import AutoModel
+
+# https://github.com/HazyResearch/hyena-dna
+
+
+class TokenClassificationModule(nn.Module):
+    """Token classification model."""
+
+    def __init__(
+        self,
+        number_of_classes: int,
+        head: nn.Module,
+        backbone_name: str = "hyenadna-small-32k-seqlen",
+        *,
+        freeze_backbone=False,
+    ):
+        super().__init__()
+        self.number_of_classes = number_of_classes
+        self.backbone_name = backbone_name
+        self.backbone = AutoModel.from_pretrained(f"LongSafari/{backbone_name}-hf", trust_remote_code=True)
+        self.head = head
+
+        if freeze_backbone:
+            for param in self.backbone.parameters():
+                param.requires_grad = False
+
+    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)

deepchopper/models/llm/metric.py

@@ -0,0 +1,44 @@
+import evaluate
+import numpy as np
+
+clf_metrics = evaluate.combine(["accuracy", "f1", "precision", "recall"])
+
+IGNORE_INDEX = -100
+
+
+def compute_metrics(p):
+    """Compute metrics for a given set of predictions and labels.
+
+    Parameters:
+    p (tuple): A tuple containing two numpy arrays - predictions and labels.
+               predictions: 3D numpy array of shape (batch_size, sequence_length, num_classes)
+               labels: 2D numpy array of shape (batch_size, sequence_length)
+
+    Returns:
+    dict: A dictionary containing computed metrics for the predictions.
+
+    Raises:
+    ValueError: If the input arrays are not of the expected shape.
+    """
+    predictions, labels = p
+    predictions = np.argmax(predictions, axis=2)
+    # Initialize lists to hold the filtered predictions and labels
+    true_predictions = []
+    true_labels = []
+
+    # Filter out '-100' labels and correspondingly filter predictions
+    for prediction, label in zip(predictions, labels, strict=True):
+        filtered_prediction = []
+        filtered_label = []
+
+        for p, l in zip(prediction, label, strict=True):
+            if l != IGNORE_INDEX:
+                filtered_prediction.append(p)
+                filtered_label.append(l)
+        true_predictions.append(filtered_prediction)
+        true_labels.append(filtered_label)
+
+    for preds, refs in zip(true_predictions, true_labels, strict=True):
+        clf_metrics.add_batch(predictions=preds, references=refs)
+
+    return clf_metrics.compute()

deepchopper/models/llm/tokenizer.py

@@ -0,0 +1,205 @@
+from functools import partial
+
+import torch
+from datasets import Dataset
+from transformers import (
+    AutoTokenizer,
+    DataCollatorForTokenClassification,
+)
+
+import deepchopper
+
+from .metric import IGNORE_INDEX
+
+
+def pad_without_fast_tokenizer_warning(tokenizer, *pad_args, **pad_kwargs):
+    """Pads without triggering the warning about how using the pad function is sub-optimal when using a fast tokenizer."""
+    # To avoid errors when using Feature extractors
+    if not hasattr(tokenizer, "deprecation_warnings"):
+        return tokenizer.pad(*pad_args, **pad_kwargs)
+
+    # Save the state of the warning, then disable it
+    warning_state = tokenizer.deprecation_warnings.get("Asking-to-pad-a-fast-tokenizer", False)
+    tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = True
+
+    try:
+        padded = tokenizer.pad(*pad_args, **pad_kwargs)
+    finally:
+        # Restore the state of the warning.
+        tokenizer.deprecation_warnings["Asking-to-pad-a-fast-tokenizer"] = warning_state
+
+    return padded
+
+
+class DataCollatorForTokenClassificationWithQual(DataCollatorForTokenClassification):
+    def torch_call(self, features):
+        import torch
+
+        label_name = "label" if "label" in features[0] else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0] else None
+
+        qual_name = "input_quals"
+        qual_pad_token_id = 0
+        input_quals = [feature[qual_name] for feature in features]
+
+        id_name = "id"  # for predction dataset
+
+        no_labels_features = [
+            {k: v for k, v in feature.items() if k not in [qual_name, label_name, id_name]} for feature in features
+        ]
+
+        batch = pad_without_fast_tokenizer_warning(
+            self.tokenizer,
+            no_labels_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+        if labels is None:
+            return batch
+
+        sequence_length = batch["input_ids"].shape[1]
+        padding_side = self.tokenizer.padding_side
+
+        def to_list(tensor_or_iterable):
+            if isinstance(tensor_or_iterable, torch.Tensor):
+                return tensor_or_iterable.tolist()
+            return list(tensor_or_iterable)
+
+        if padding_side == "right":
+            batch[label_name] = [
+                to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels
+            ]
+            batch[qual_name] = [
+                to_list(qual) + [qual_pad_token_id] * (sequence_length - len(qual)) for qual in input_quals
+            ]
+        else:
+            batch[label_name] = [
+                [self.label_pad_token_id] * (sequence_length - len(label)) + to_list(label) for label in labels
+            ]
+            batch[qual_name] = [
+                [qual_pad_token_id] * (sequence_length - len(qual)) + to_list(qual) for qual in input_quals
+            ]
+
+        batch[label_name] = torch.tensor(batch[label_name], dtype=torch.int8)
+        batch[qual_name] = torch.tensor(batch[qual_name], dtype=torch.float32)
+
+        # for predction dataset and save id feature
+        if id_name in features[0]:
+            batch[id_name] = torch.tensor([to_list(feature[id_name]) for feature in features], dtype=torch.int8)
+
+        return batch
+
+
+def load_tokenizer_from_hyena_model(model_name):
+    max_lengths = {
+        "hyenadna-tiny-1k-seqlen": 1024,
+        "hyenadna-small-32k-seqlen": 32768,
+        "hyenadna-medium-160k-seqlen": 160000,
+        "hyenadna-medium-450k-seqlen": 450000,  # T4 up to here
+        "hyenadna-large-1m-seqlen": 1_000_000,  # only A100 (paid tier)
+    }
+
+    if model_name not in max_lengths:
+        msg = f"Model name {model_name} not found in available models."
+        raise ValueError(msg)
+
+    max_length = max_lengths[model_name]
+    # bfloat16 for better speed and reduced memory usage
+    model_name = f"LongSafari/{model_name}-hf"
+    return AutoTokenizer.from_pretrained(
+        model_name, max_length=max_length, truncation=True, padding=True, trust_remote_code=True, force_download=False
+    )
+
+
+def to_list(data):
+    return list(data)
+
+
+def tokenize_and_align_labels_and_quals(data, tokenizer, max_length, pad_qual=0, pad_label=IGNORE_INDEX):
+    tokenized_inputs = tokenizer(data["seq"], max_length=max_length, truncation=True, padding=True)
+
+    if isinstance(data["qual"], bytes):
+        data["qual"] = torch.Tensor(to_list(data["qual"]))
+
+    if len(data["seq"]) >= max_length:
+        if data["target"][1] + 2 > max_length:
+            labels = torch.tensor([*deepchopper.vectorize_target(0, 0, max_length - 1), pad_label])
+            quals = torch.cat((data["qual"][: max_length - 1], torch.tensor([pad_qual]))).float()
+            normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+        else:
+            labels = torch.tensor([*deepchopper.vectorize_targets(data["target"].numpy(), max_length - 1), pad_label])
+            quals = torch.cat((data["qual"][: max_length - 1], torch.tensor([pad_qual]))).float()
+            normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+    else:
+        labels = torch.tensor([*deepchopper.vectorize_targets(data["target"].numpy(), len(data["seq"])), pad_label])
+        quals = torch.cat((data["qual"], torch.tensor([pad_qual]))).float()
+        normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+
+    tokenized_inputs.update({"labels": labels, "input_quals": normalized_quals})
+    return tokenized_inputs
+
+
+def tokenize_and_align_labels_and_quals_ids(
+    data, tokenizer, max_length, pad_qual=0, pad_label=IGNORE_INDEX, max_id_length=256
+):
+    tokenized_inputs = tokenizer(data["seq"], max_length=max_length, truncation=True, padding=True)
+
+    truncation = False
+
+    # TODO: remove target and labels during prediction
+
+    if len(data["seq"]) >= max_length:
+        truncation = True
+        if data["target"][1] + 2 > max_length:
+            labels = torch.tensor([*deepchopper.vectorize_target(0, 0, max_length - 1), pad_label])
+            quals = torch.cat((data["qual"][: max_length - 1], torch.tensor([pad_qual]))).float()
+            normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+        else:
+            labels = torch.tensor([*deepchopper.vectorize_targets(data["target"].numpy(), max_length - 1), pad_label])
+            quals = torch.cat((data["qual"][: max_length - 1], torch.tensor([pad_qual]))).float()
+            normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+    else:
+        labels = torch.tensor([*deepchopper.vectorize_targets(data["target"].numpy(), len(data["seq"])), pad_label])
+        quals = torch.cat((data["qual"], torch.tensor([pad_qual]))).float()
+        normalized_quals = torch.nn.functional.normalize(quals, dim=0)
+
+    # change id to ascii values
+    new_id = [len(data["id"]), int(truncation)]
+    new_id += [ord(char) for char in data["id"]]
+    if len(new_id) > max_id_length:
+        new_id = new_id[:max_id_length]
+    elif len(new_id) < max_id_length:
+        new_id += [0] * (max_id_length - len(new_id))
+
+    tokenized_inputs.update({"labels": labels, "input_quals": normalized_quals, "id": new_id})
+    return tokenized_inputs
+
+
+def tokenize_dataset(dataset, tokenizer, max_length):
+    """Tokenizes the input dataset using the provided tokenizer and aligns labels and qualities.
+
+    Args:
+        dataset (Dataset): The input dataset to be tokenized.
+        tokenizer (Tokenizer): The tokenizer to be used for tokenization.
+        max_length (int): The maximum length of the tokenized sequences.
+
+    Returns:
+        Tokenized dataset with aligned labels and qualities.
+
+    Raises:
+        ValueError: If the dataset is empty or if the tokenizer is not provided.
+        TypeError: If the dataset is not of type Dataset or if the tokenizer is not of type Tokenizer.
+    """
+    if not dataset:
+        raise ValueError("Input dataset is empty")
+    if not tokenizer:
+        raise ValueError("Tokenizer is not provided")
+    if not isinstance(dataset, Dataset):
+        raise TypeError("Input dataset must be of type Dataset")
+
+    return dataset.map(
+        partial(tokenize_and_align_labels_and_quals, tokenizer=tokenizer, max_length=max_length)
+    ).remove_columns(["id", "seq", "qual", "target"])

deepchopper/models/transformer.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import nn
+from torch.nn import (
+    TransformerDecoder,
+    TransformerDecoderLayer,
+    TransformerEncoder,
+    TransformerEncoderLayer,
+)
+
+
+class TokenClassificationModule(nn.Module):
+    def __init__(self, vocab_size, d_model, nhead, num_encoder_layers, dim_feedforward, number_of_classes):
+        super().__init__()
+        self.number_of_classes = number_of_classes
+
+        self.qual_linear1 = nn.Sequential(
+            nn.Linear(1, d_model),
+        )
+
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        encoder_layer = TransformerEncoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
+        )
+
+        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
+
+        self.classifier = nn.Linear(d_model, number_of_classes)
+        self.d_model = d_model
+
+    def forward(self, src: torch.Tensor, input_quals: torch.Tensor, src_mask=None):
+        seq_length = src.size(1)
+
+        # Calculate positional embeddings
+        position = torch.arange(seq_length, dtype=torch.float, device=src.device).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, self.d_model, 2, device=src.device).float()
+            * (-torch.log(torch.tensor(10000.0)) / self.d_model)
+        )
+
+        pos_embedding = torch.zeros((seq_length, self.d_model), device=src.device)
+        pos_embedding[:, 0::2] = torch.sin(position * div_term)
+        pos_embedding[:, 1::2] = torch.cos(position * div_term)
+
+        src = F.relu(self.embedding(src) + self.qual_linear1(input_quals.unsqueeze(-1)))
+        src = src + pos_embedding.unsqueeze(0)
+
+        if src_mask is not None:
+            src_mask = src_mask.to(dtype=torch.bool)
+
+        output = self.transformer_encoder(src, src_key_padding_mask=src_mask)
+        return self.classifier(output)
+
+
+class Seq2SeqTokenClassifier(nn.Module):
+    def __init__(
+        self,
+        vocab_size,
+        d_model,
+        nhead,
+        num_encoder_layers,
+        num_decoder_layers,
+        dim_feedforward,
+        number_of_classes,
+    ):
+        super().__init__()
+        self.d_model = d_model
+
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.qual_linear1 = nn.Linear(1, d_model)
+
+        encoder_layer = TransformerEncoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
+        )
+        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
+
+        decoder_layer = TransformerDecoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, batch_first=True
+        )
+        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
+
+        self.classifier = nn.Linear(d_model, number_of_classes)
+
+    def forward(
+        self,
+        src: torch.Tensor,
+        input_quals: torch.Tensor,
+        tgt: torch.Tensor,
+        src_mask=None,
+        tgt_mask=None,
+    ):
+        src_embeddings = self.embedding(src)
+        qual_embeddings = self.qual_linear1(input_quals.unsqueeze(-1))
+
+        encoder_input = src_embeddings + qual_embeddings
+
+        if src_mask is not None:
+            src_mask = src_mask.to(dtype=torch.bool)
+
+        encoder_output = self.transformer_encoder(encoder_input, src_key_padding_mask=src_mask)
+
+        tgt_embeddings = self.embedding(tgt)
+        decoder_output = self.transformer_decoder(
+            tgt_embeddings, encoder_output, memory_key_padding_mask=src_mask, tgt_mask=tgt_mask
+        )
+
+        return self.classifier(decoder_output)