omnigenome 0.3.0a1__py3-none-any.whl → 0.3.3a0__py3-none-any.whl

omnigenome/src/model/embedding/model.py

@@ -1,263 +0,0 @@
-# -*- coding: utf-8 -*-
-# file: model.py
-# time: 18:37 22/09/2024
-# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
-# github: https://github.com/yangheng95
-# huggingface: https://huggingface.co/yangheng
-# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
-# Copyright (C) 2019-2024. All Rights Reserved.
-
-import torch
-from transformers import AutoTokenizer, AutoModel
-
-from omnigenome.src.misc.utils import fprint
-
-
-class OmniModelForEmbedding(torch.nn.Module):
-    """
-    A wrapper class for generating embeddings from pre-trained models.
-    
-    This class provides a unified interface for loading pre-trained models and
-    generating embeddings from genomic sequences. It supports various aggregation
-    methods and batch processing for efficient embedding generation.
-    
-    Attributes:
-        tokenizer: The tokenizer for processing input sequences
-        model: The pre-trained model for generating embeddings
-        _device: The device (CPU/GPU) where the model is loaded
-        
-    Example:
-        >>> from omnigenome import OmniModelForEmbedding
-        >>> model = OmniModelForEmbedding("anonymous8/OmniGenome-186M")
-        >>> sequences = ["ATCGGCTA", "GGCTAGCTA"]
-        >>> embeddings = model.batch_encode(sequences)
-        >>> print(f"Embeddings shape: {embeddings.shape}")
-        torch.Size([2, 768])
-    """
-    
-    def __init__(self, model_name_or_path, *args, **kwargs):
-        """
-        Initialize the embedding model.
-        
-        Args:
-            model_name_or_path (str): Name or path of the pre-trained model to load
-            *args: Additional positional arguments passed to AutoModel.from_pretrained
-            **kwargs: Additional keyword arguments passed to AutoModel.from_pretrained
-        """
-        super().__init__()
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
-        self.model = AutoModel.from_pretrained(model_name_or_path, *args, **kwargs)
-        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self._device)
-        self.model.eval()  # Set model to evaluation mode
-
-    def batch_encode(self, sequences, batch_size=8, max_length=512, agg='head'):
-        """
-        Encode a list of sequences to their corresponding embeddings.
-        
-        This method processes sequences in batches for memory efficiency and
-        supports different aggregation methods for the final embeddings.
-        
-        Args:
-            sequences (list of str): List of input sequences to encode
-            batch_size (int, optional): Batch size for processing. Defaults to 8
-            max_length (int, optional): Maximum sequence length for encoding. Defaults to 512
-            agg (str, optional): Aggregation method for embeddings. Options are 'head', 'mean', 'tail'. Defaults to 'head'
-            
-        Returns:
-            torch.Tensor: Embeddings for the input sequences with shape (n_sequences, embedding_dim)
-            
-        Raises:
-            ValueError: If unsupported aggregation method is provided
-            
-        Example:
-            >>> sequences = ["ATCGGCTA", "GGCTAGCTA", "TATCGCTA"]
-            >>> embeddings = model.batch_encode(sequences, batch_size=2, agg='mean')
-            >>> print(f"Embeddings shape: {embeddings.shape}")
-            torch.Size([3, 768])
-        """
-        embeddings = []
-
-        for i in range(0, len(sequences), batch_size):
-            batch_sequences = sequences[i: i + batch_size]
-            inputs = self.tokenizer(
-                batch_sequences,
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=max_length,
-            )
-            inputs = {key: value.to(self.device) for key, value in inputs.items()}
-
-            with torch.no_grad():
-                outputs = self.model(**inputs)
-
-            batch_embeddings = outputs.last_hidden_state.cpu()
-
-            if agg == 'head':
-                emb = batch_embeddings[:, 0, :]
-            elif agg == 'mean':
-                attention_mask = inputs["attention_mask"].cpu()
-                masked_embeddings = batch_embeddings * attention_mask.unsqueeze(-1)
-                lengths = attention_mask.sum(dim=1).unsqueeze(1)
-                emb = masked_embeddings.sum(dim=1) / lengths
-            elif agg == 'tail':
-                attention_mask = inputs["attention_mask"]
-                lengths = attention_mask.sum(dim=1) - 1
-                emb = torch.stack([
-                    batch_embeddings[i, l.item()] for i, l in enumerate(lengths)
-                ])
-            else:
-                raise ValueError(f"Unsupported aggregation method: {agg}")
-
-            embeddings.append(emb)
-
-        embeddings = torch.cat(embeddings, dim=0)
-        fprint(f"Generated embeddings for {len(sequences)} sequences.")
-        return embeddings
-
-    def encode(self, sequence, max_length=512, agg='head', keep_dim=False):
-        """
-        Encode a single sequence to its corresponding embedding.
-        
-        Args:
-            sequence (str): Input sequence to encode
-            max_length (int, optional): Maximum sequence length for encoding. Defaults to 512
-            agg (str, optional): Aggregation method. Options are 'head', 'mean', 'tail'. Defaults to 'head'
-            keep_dim (bool, optional): Whether to retain the batch dimension. Defaults to False
-            
-        Returns:
-            torch.Tensor: Embedding for the input sequence
-            
-        Raises:
-            ValueError: If unsupported aggregation method is provided
-            
-        Example:
-            >>> sequence = "ATCGGCTA"
-            >>> embedding = model.encode(sequence, agg='mean')
-            >>> print(f"Embedding shape: {embedding.shape}")
-            torch.Size([768])
-        """
-        inputs = self.tokenizer(
-            sequence,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=max_length,
-        )
-        inputs = {key: value.to(self.device) for key, value in inputs.items()}
-
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-
-        last_hidden = outputs.last_hidden_state.cpu()
-
-        if agg == 'head':
-            emb = last_hidden[0, 0]
-        elif agg == 'mean':
-            attention_mask = inputs["attention_mask"].cpu()
-            masked_embeddings = last_hidden * attention_mask.unsqueeze(-1)
-            lengths = attention_mask.sum(dim=1).unsqueeze(1)
-            emb = masked_embeddings.sum(dim=1) / lengths
-            emb = emb.squeeze(0)
-        elif agg == 'tail':
-            attention_mask = inputs["attention_mask"]
-            lengths = attention_mask.sum(dim=1) - 1
-            emb = last_hidden[0, lengths[0].item()]
-        else:
-            raise ValueError(f"Unsupported aggregation method: {agg}")
-
-        return emb.unsqueeze(0) if keep_dim else emb
-
-    def save_embeddings(self, embeddings, output_path):
-        """
-        Save the generated embeddings to a file.
-        
-        Args:
-            embeddings (torch.Tensor): The embeddings to save
-            output_path (str): Path to save the embeddings
-            
-        Example:
-            >>> embeddings = model.batch_encode(sequences)
-            >>> model.save_embeddings(embeddings, "embeddings.pt")
-            >>> print("Embeddings saved successfully")
-        """
-        torch.save(embeddings, output_path)
-        fprint(f"Embeddings saved to {output_path}")
-
-    def load_embeddings(self, embedding_path):
-        """
-        Load embeddings from a file.
-        
-        Args:
-            embedding_path (str): Path to the saved embeddings
-            
-        Returns:
-            torch.Tensor: The loaded embeddings
-            
-        Example:
-            >>> embeddings = model.load_embeddings("embeddings.pt")
-            >>> print(f"Loaded embeddings shape: {embeddings.shape}")
-            torch.Size([100, 768])
-        """
-        embeddings = torch.load(embedding_path)
-        fprint(f"Loaded embeddings from {embedding_path}")
-        return embeddings
-
-    def compute_similarity(self, embedding1, embedding2, dim=0):
-        """
-        Compute cosine similarity between two embeddings.
-        
-        Args:
-            embedding1 (torch.Tensor): The first embedding
-            embedding2 (torch.Tensor): The second embedding
-            dim (int, optional): Dimension along which to compute cosine similarity. Defaults to 0
-            
-        Returns:
-            float: Cosine similarity score between -1 and 1
-            
-        Example:
-            >>> emb1 = model.encode("ATCGGCTA")
-            >>> emb2 = model.encode("GGCTAGCTA")
-            >>> similarity = model.compute_similarity(emb1, emb2)
-            >>> print(f"Cosine similarity: {similarity:.4f}")
-            0.8234
-        """
-        similarity = torch.nn.functional.cosine_similarity(
-            embedding1, embedding2, dim=dim
-        )
-        return similarity
-
-    @property
-    def device(self):
-        """
-        Get the current device ('cuda' or 'cpu').
-        
-        Returns:
-            torch.device: The device where the model is loaded
-        """
-        return self._device
-
-
-# Example usage
-if __name__ == "__main__":
-    model_name = "anonymous8/OmniGenome-186M"
-    embedding_model = OmniModelForEmbedding(model_name)
-
-    # Encode multiple sequences
-    sequences = ["ATCGGCTA", "GGCTAGCTA"]
-    embedding = embedding_model.encode(sequences[0])
-    fprint(f"Single embedding shape: {embedding.shape}")
-
-    embeddings = embedding_model.batch_encode(sequences)
-    fprint(f"Embeddings for sequences: {embeddings}")
-
-    # Save and load embeddings
-    embedding_model.save_embeddings(embeddings, "embeddings.pt")
-    loaded_embeddings = embedding_model.load_embeddings("embeddings.pt")
-
-    # Compute similarity between two embeddings
-    similarity = embedding_model.compute_similarity(
-        loaded_embeddings[0], loaded_embeddings[1]
-    )
-    fprint(f"Cosine similarity: {similarity}")

omnigenome/src/model/mlm/__init__.py

@@ -1,12 +0,0 @@
-# -*- coding: utf-8 -*-
-# file: __init__.py
-# time: 13:30 10/04/2024
-# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
-# github: https://github.com/yangheng95
-# huggingface: https://huggingface.co/yangheng
-# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
-# Copyright (C) 2019-2024. All Rights Reserved.
-"""
-This package contains modules for Masked Language Models (MLM).
-"""
-

omnigenome/src/model/mlm/model.py

@@ -1,177 +0,0 @@
-# -*- coding: utf-8 -*-
-# file: model.py
-# time: 13:30 10/04/2024
-# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
-# github: https://github.com/yangheng95
-# huggingface: https://huggingface.co/yangheng
-# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
-# Copyright (C) 2019-2024. All Rights Reserved.
-"""
-Masked Language Model (MLM) for genomic sequences.
-
-This module provides a masked language model implementation specifically designed
-for genomic sequences. It supports masked language modeling tasks where tokens
-are randomly masked and the model learns to predict the original tokens.
-"""
-import numpy as np
-import torch
-from transformers import BatchEncoding
-
-from ...abc.abstract_model import OmniModel
-
-
-class OmniModelForMLM(OmniModel):
-    """
-    Masked Language Model for genomic sequences.
-    
-    This model implements masked language modeling for genomic sequences, where
-    tokens are randomly masked and the model learns to predict the original tokens.
-    It's useful for pre-training genomic language models and understanding sequence
-    patterns and dependencies.
-    
-    Attributes:
-        loss_fn: Cross-entropy loss function for masked language modeling
-    """
-    
-    def __init__(self, config_or_model, tokenizer, *args, **kwargs):
-        """
-        Initialize the MLM model.
-        
-        Args:
-            config_or_model: Model configuration or pre-trained model
-            tokenizer: Tokenizer for processing input sequences
-            *args: Additional positional arguments
-            **kwargs: Additional keyword arguments
-            
-        Raises:
-            ValueError: If the model doesn't support masked language modeling
-        """
-        super().__init__(config_or_model, tokenizer, *args, **kwargs)
-        self.metadata["model_name"] = self.__class__.__name__
-        if "MaskedLM" not in self.model.__class__.__name__:
-            raise ValueError(
-                "The model does not have a language model head, which is required for MLM."
-                "Please use a model that supports masked language modeling."
-            )
-
-        self.loss_fn = torch.nn.CrossEntropyLoss()
-
-    def forward(self, **inputs):
-        """
-        Forward pass for masked language modeling.
-        
-        Args:
-            **inputs: Input tensors including input_ids, attention_mask, and labels
-            
-        Returns:
-            dict: Dictionary containing loss, logits, and last_hidden_state
-        """
-        inputs = inputs.pop("inputs")
-        outputs = self.model(**inputs, output_hidden_states=True)
-        last_hidden_state = (
-            outputs["last_hidden_state"]
-            if "last_hidden_state" in outputs
-            else outputs["hidden_states"][-1]
-        )
-        logits = outputs["logits"] if "logits" in outputs else None
-        loss = outputs["loss"] if "loss" in outputs else None
-        outputs = {
-            "loss": loss,
-            "logits": logits,
-            "last_hidden_state": last_hidden_state,
-        }
-        return outputs
-
-    def predict(self, sequence_or_inputs, **kwargs):
-        """
-        Generate predictions for masked language modeling.
-        
-        Args:
-            sequence_or_inputs: Input sequences or pre-processed inputs
-            **kwargs: Additional keyword arguments
-            
-        Returns:
-            dict: Dictionary containing predictions, logits, and last_hidden_state
-        """
-        raw_outputs = self._forward_from_raw_input(sequence_or_inputs, **kwargs)
-
-        logits = raw_outputs["logits"]
-        last_hidden_state = raw_outputs["last_hidden_state"]
-
-        predictions = []
-        for i in range(logits.shape[0]):
-            predictions.append(logits[i].argmax(dim=-1).cpu())
-
-        if not isinstance(sequence_or_inputs, list):
-            outputs = {
-                "predictions": predictions[0],
-                "logits": logits[0],
-                "last_hidden_state": last_hidden_state[0],
-            }
-        else:
-            outputs = {
-                "predictions": (
-                    torch.stack(predictions)
-                    if predictions[0].shape
-                    else torch.tensor(predictions).to(self.model.device)
-                ),
-                "logits": logits,
-                "last_hidden_state": last_hidden_state,
-            }
-
-        return outputs
-
-    def inference(self, sequence_or_inputs, **kwargs):
-        """
-        Perform inference for masked language modeling, decoding predictions to sequences.
-        
-        Args:
-            sequence_or_inputs: Input sequences or pre-processed inputs
-            **kwargs: Additional keyword arguments
-            
-        Returns:
-            dict: Dictionary containing decoded predictions, logits, and last_hidden_state
-        """
-        raw_outputs = self._forward_from_raw_input(sequence_or_inputs, **kwargs)
-
-        inputs = raw_outputs["inputs"]
-        logits = raw_outputs["logits"]
-        last_hidden_state = raw_outputs["last_hidden_state"]
-
-        predictions = []
-        for i in range(logits.shape[0]):
-            i_logit = logits[i][inputs["input_ids"][i].ne(self.config.pad_token_id)][
-                1:-1
-            ]
-            prediction = self.tokenizer.decode(i_logit.argmax(dim=-1)).replace(" ", "")
-            predictions.append(list(prediction))
-
-        if not isinstance(sequence_or_inputs, list):
-            outputs = {
-                "predictions": predictions[0],
-                "logits": logits[0],
-                "last_hidden_state": last_hidden_state[0],
-            }
-        else:
-            outputs = {
-                "predictions": predictions,
-                "logits": logits,
-                "last_hidden_state": last_hidden_state,
-            }
-
-        return outputs
-
-    def loss_function(self, logits, labels):
-        """
-        Compute the loss for masked language modeling.
-        
-        Args:
-            logits (torch.Tensor): Model predictions [batch_size, seq_len, vocab_size]
-            labels (torch.Tensor): Ground truth labels [batch_size, seq_len]
-            
-        Returns:
-            torch.Tensor: Computed cross-entropy loss value
-        """
-        loss_fn = torch.nn.CrossEntropyLoss()
-        loss = loss_fn(logits.view(-1, self.tokenizer.vocab_size), labels.view(-1))
-        return loss

omnigenome/src/model/module_utils.py

@@ -1,232 +0,0 @@
-# -*- coding: utf-8 -*-
-# file: module_utils.py
-# time: 22:53 18/07/2024
-# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
-# github: https://github.com/yangheng95
-# huggingface: https://huggingface.co/yangheng
-# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
-# Copyright (C) 2019-2024. All Rights Reserved.
-"""
-Module utilities for OmniGenome models.
-
-This module provides utility classes and functions for handling model inputs,
-pooling operations, and attention mechanisms used across different OmniGenome model types.
-"""
-import torch
-import torch.nn as nn
-
-from transformers.models.bert.modeling_bert import BertPooler
-from transformers.tokenization_utils_base import BatchEncoding
-
-
-class OmniPooling(torch.nn.Module):
-    """
-    A flexible pooling layer for OmniGenome models that handles different input formats.
-    
-    This class provides a unified interface for pooling operations across different
-    model architectures, supporting both causal language models and encoder-based models.
-    It can handle various input formats including tuples, dictionaries, BatchEncoding
-    objects, and tensors.
-    
-    Attributes:
-        config: Model configuration object containing architecture and tokenizer settings
-        pooler: BertPooler instance for non-causal models, None for causal models
-    """
-    
-    def __init__(self, config, *args, **kwargs):
-        """
-        Initialize the OmniPooling layer.
-        
-        Args:
-            config: Model configuration object containing architecture information
-            *args: Additional positional arguments
-            **kwargs: Additional keyword arguments
-        """
-        super().__init__(*args, **kwargs)
-        self.config = config
-        self.pooler = BertPooler(self.config) if not self._is_causal_lm() else None
-
-    def forward(self, inputs, last_hidden_state):
-        """
-        Perform pooling operation on the last hidden state.
-        
-        This method handles different input formats and applies appropriate pooling:
-        - For causal language models: Uses the last non-padded token
-        - For encoder models: Uses the BertPooler
-        
-        Args:
-            inputs: Input data in various formats (tuple, dict, BatchEncoding, or tensor)
-            last_hidden_state (torch.Tensor): Hidden states from the model [batch_size, seq_len, hidden_size]
-            
-        Returns:
-            torch.Tensor: Pooled representation [batch_size, hidden_size]
-            
-        Raises:
-            ValueError: If input format is not supported or cannot be parsed
-        """
-        if isinstance(inputs, tuple):
-            input_ids = inputs[0]
-            attention_mask = inputs[1] if len(inputs) > 1 else None
-        elif isinstance(inputs, BatchEncoding) or isinstance(inputs, dict):
-            input_ids = inputs["input_ids"]
-            attention_mask = (
-                inputs["attention_mask"] if "attention_mask" in inputs else None
-            )
-        elif isinstance(inputs, torch.Tensor):
-            shape = inputs.shape
-            try:
-                if len(shape) == 3:
-                    # compatible with hf_trainer in AutoBenchmark
-                    if shape[1] == 2:
-                        input_ids = inputs[:, 0]
-                        attention_mask = inputs[:, 1]
-                    else:
-                        input_ids = inputs[0]
-                        attention_mask = inputs[1] if len(inputs) > 1 else None
-                elif len(shape) == 2:
-                    input_ids = inputs
-                    attention_mask = None
-            except:
-                raise ValueError(
-                    f"Failed to get the input_ids and attention_mask from the inputs, got shape {shape}."
-                )
-        else:
-            raise ValueError(
-                f"The inputs should be a tuple, BatchEncoding or a dictionary-like object, got {type(inputs)}."
-            )
-
-        if not self.pooler:
-            pad_token_id = getattr(self.config, "pad_token_id", -100)
-            sequence_lengths = input_ids.ne(pad_token_id).sum(dim=1) - 1
-            last_hidden_state = last_hidden_state[
-                torch.arange(input_ids.size(0), device=last_hidden_state.device),
-                sequence_lengths,
-            ]
-        else:
-            last_hidden_state = self.pooler(last_hidden_state)
-
-        return last_hidden_state
-
-    def _is_causal_lm(self):
-        """
-        Check if the model is a causal language model.
-        
-        Determines if the model architecture is causal based on the configuration.
-        
-        Returns:
-            bool: True if the model is a causal language model, False otherwise
-        """
-        if (
-            hasattr(self.config, "architectures")
-            and "CausalLM" in str(self.config.architectures)
-        ) or (
-            hasattr(self.config, "auto_map") and "CausalLM" in str(self.config.auto_map)
-        ):
-            return True
-        else:
-            return False
-
-
-# class InteractingAttention(nn.Module):
-#     def __init__(self, embed_size, num_heads=12):
-#         super(InteractingAttention, self).__init__()
-#         self.num_heads = num_heads
-#         self.embed_size = embed_size
-#
-#         assert embed_size % num_heads == 0, "Embedding size should be divisible by number of heads"
-#
-#         self.head_dim = embed_size // num_heads
-#
-#         self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
-#         self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
-#         self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
-#         self.layer_norm = nn.LayerNorm(num_heads * self.head_dim, eps=1e-6)
-#
-#         self.fc_out = nn.Linear(num_heads * self.head_dim, embed_size)
-#
-#     # def forward(self, query, keys, values):
-#     def forward(self, query, keys, values):
-#
-#         N = query.shape[0]
-#         value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]
-#
-#         # Split embedding into self.num_heads pieces
-#         values = values.reshape(N, value_len, self.num_heads, self.head_dim)
-#         keys = keys.reshape(N, key_len, self.num_heads, self.head_dim)
-#         queries = query.reshape(N, query_len, self.num_heads, self.head_dim)
-#
-#         values = self.values(values)  # (N, value_len, heads, head_dim)
-#         keys = self.keys(keys)  # (N, key_len, heads, head_dim)
-#         queries = self.queries(queries)  # (N, query_len, heads, head_dim)
-#
-#         energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
-#
-#         attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)  # (N, heads, query_len, key_len)
-#
-#         out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(
-#             N, query_len, self.num_heads * self.head_dim
-#         )
-#         out = self.layer_norm(out + query)
-#         out = self.fc_out(out)
-#         out = self.layer_norm(out + query)
-#         return out
-
-
-class InteractingAttention(nn.Module):
-    """
-    An interacting attention mechanism for sequence modeling.
-    
-    This class implements a multi-head attention mechanism with residual connections
-    and layer normalization. It's designed for processing sequences where different
-    parts of the sequence need to interact with each other.
-    
-    Attributes:
-        attention: Multi-head attention layer
-        layer_norm: Layer normalization for residual connections
-        fc_out: Output projection layer
-    """
-    
-    def __init__(self, embed_size, num_heads=24):
-        """
-        Initialize the InteractingAttention module.
-        
-        Args:
-            embed_size (int): Size of the embedding dimension
-            num_heads (int): Number of attention heads (default: 24)
-            
-        Raises:
-            AssertionError: If embed_size is not divisible by num_heads
-        """
-        super(InteractingAttention, self).__init__()
-        assert (
-            embed_size % num_heads == 0
-        ), "Embedding size should be divisible by number of heads"
-
-        self.attention = nn.MultiheadAttention(
-            embed_dim=embed_size, num_heads=num_heads, batch_first=True
-        )
-
-        self.layer_norm = nn.LayerNorm(embed_size, eps=1e-6)
-
-        self.fc_out = nn.Linear(embed_size, embed_size)
-
-    def forward(self, query, keys, values):
-        """
-        Forward pass through the interacting attention mechanism.
-        
-        Args:
-            query (torch.Tensor): Query tensor [batch_size, query_len, embed_size]
-            keys (torch.Tensor): Key tensor [batch_size, key_len, embed_size]
-            values (torch.Tensor): Value tensor [batch_size, value_len, embed_size]
-            
-        Returns:
-            torch.Tensor: Output tensor with same shape as query
-        """
-        att_output, _ = self.attention(query, keys, values)
-
-        query = self.layer_norm(att_output + query)
-
-        output = self.fc_out(query)
-        output = self.layer_norm(output + query)
-
-        return output

omnigenome/src/model/regression/__init__.py

@@ -1,12 +0,0 @@
-# -*- coding: utf-8 -*-
-# file: __init__.py
-# time: 21:11 08/04/2024
-# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
-# github: https://github.com/yangheng95
-# huggingface: https://huggingface.co/yangheng
-# google scholar: https://scholar.google.com/citations?user=NPq5a_0AAAAJ&hl=en
-# Copyright (C) 2019-2024. All Rights Reserved.
-"""
-This package contains modules for regression models.
-"""
-