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

omnigenome/src/abc/abstract_model.py

@@ -47,14 +47,14 @@ def count_parameters(model):
 class OmniModel(torch.nn.Module):
     """
     Abstract base class for all models in OmniGenome.
-    
+
     This class provides a unified interface for all genomic models in the OmniGenome
     framework. It handles model initialization, forward passes, loss computation,
     prediction, inference, and model persistence.
-    
+
     The class is designed to work with various types of genomic data and tasks,
     including sequence classification, token classification, regression, and more.
-    
+
     Attributes:
         model (torch.nn.Module): The underlying PyTorch model.
         config: The model configuration.
@@ -76,16 +76,16 @@ class OmniModel(torch.nn.Module):
         - From a configuration object
 
         Args:
-            config_or_model: A model configuration, a pre-trained model path (str), 
+            config_or_model: A model configuration, a pre-trained model path (str),
                            or a `torch.nn.Module` instance.
             tokenizer: The tokenizer associated with the model.
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
                 - label2id (dict): Mapping from class labels to IDs.
                 - num_labels (int): The number of labels.
-                - trust_remote_code (bool): Whether to trust remote code when loading 
+                - trust_remote_code (bool): Whether to trust remote code when loading
                   from Hugging Face Hub. Defaults to True.
-                - ignore_mismatched_sizes (bool): Whether to ignore size mismatches 
+                - ignore_mismatched_sizes (bool): Whether to ignore size mismatches
                   when loading pre-trained weights. Defaults to False.
                 - dropout (float): Dropout rate. Defaults to 0.0.
 
@@ -97,7 +97,7 @@ class OmniModel(torch.nn.Module):
         Example:
             >>> # Initialize from a pre-trained model
             >>> model = OmniModelForSequenceClassification("model_path", tokenizer)
-            
+
             >>> # Initialize from a configuration
             >>> config = AutoConfig.from_pretrained("model_path")
             >>> model = OmniModelForSequenceClassification(config, tokenizer)
@@ -202,7 +202,9 @@ class OmniModel(torch.nn.Module):
             )
             self.config.num_labels = len(self.config.id2label)
 
-        assert len(self.config.label2id) == num_labels, f"Expected {num_labels} labels, but got {len(self.config.label2id)} in label2id dictionary."
+        assert (
+            len(self.config.label2id) == num_labels
+        ), f"Expected {num_labels} labels, but got {len(self.config.label2id)} in label2id dictionary."
 
         # The metadata of the model
         self.metadata = env_meta_info()
@@ -240,7 +242,7 @@ class OmniModel(torch.nn.Module):
         model architectures by mapping input parameters appropriately.
 
         Args:
-            **inputs: The inputs to the model, compatible with the base model's 
+            **inputs: The inputs to the model, compatible with the base model's
                      forward method. Typically includes 'input_ids', 'attention_mask',
                      and other model-specific parameters.
 
@@ -386,7 +388,7 @@ class OmniModel(torch.nn.Module):
         predictions for further processing.
 
         Args:
-            sequence_or_inputs: A sequence (str), list of sequences, or 
+            sequence_or_inputs: A sequence (str), list of sequences, or
                                tokenized inputs (dict/tuple).
             **kwargs: Additional arguments for tokenization and inference.
 
@@ -398,7 +400,7 @@ class OmniModel(torch.nn.Module):
         Example:
             >>> # Predict on a single sequence
             >>> outputs = model.predict("ATCGATCG")
-            
+
             >>> # Predict on multiple sequences
             >>> outputs = model.predict(["ATCGATCG", "GCTAGCTA"])
         """
@@ -416,7 +418,7 @@ class OmniModel(torch.nn.Module):
         to class labels or probabilities.
 
         Args:
-            sequence_or_inputs: A sequence (str), list of sequences, or 
+            sequence_or_inputs: A sequence (str), list of sequences, or
                                tokenized inputs (dict/tuple).
             **kwargs: Additional arguments for tokenization and inference.
 
@@ -429,7 +431,7 @@ class OmniModel(torch.nn.Module):
             >>> # Inference on a single sequence
             >>> results = model.inference("ATCGATCG")
             >>> print(results['predictions'])  # Class labels
-            
+
             >>> # Inference on multiple sequences
             >>> results = model.inference(["ATCGATCG", "GCTAGCTA"])
         """
@@ -686,4 +688,3 @@ class OmniModel(torch.nn.Module):
         info += f"Model Config: {self.config}\n"
         fprint(info)
         return info
-

omnigenome/src/abc/abstract_tokenizer.py

@@ -16,15 +16,15 @@ from ..misc.utils import env_meta_info, load_module_from_path
 class OmniTokenizer:
     """
     A wrapper class for tokenizers to provide a consistent interface within OmniGenome.
-    
+
     This class provides a unified interface for tokenizers in the OmniGenome framework.
     It wraps underlying tokenizers (typically from Hugging Face) and provides
     additional functionality for genomic sequence processing.
-    
+
     The class handles various tokenization strategies and provides compatibility
     with different model architectures. It also supports custom tokenizer wrappers
     for specialized genomic tasks.
-    
+
     Attributes:
         base_tokenizer: The underlying tokenizer instance (e.g., from Hugging Face).
         max_length (int): The default maximum sequence length.
@@ -52,7 +52,7 @@ class OmniTokenizer:
             >>> from transformers import AutoTokenizer
             >>> base_tokenizer = AutoTokenizer.from_pretrained("model_name")
             >>> tokenizer = OmniTokenizer(base_tokenizer, max_length=512)
-            
+
             >>> # Initialize with sequence conversion
             >>> tokenizer = OmniTokenizer(base_tokenizer, u2t=True)
         """
@@ -87,9 +87,9 @@ class OmniTokenizer:
         Example:
             >>> # Load from a pre-trained model
             >>> tokenizer = OmniTokenizer.from_pretrained("model_name")
-            
+
             >>> # Load with custom parameters
-            >>> tokenizer = OmniTokenizer.from_pretrained("model_name", 
+            >>> tokenizer = OmniTokenizer.from_pretrained("model_name",
             ...                                          trust_remote_code=True)
         """
         wrapper_path = f"{model_name_or_path.rstrip('/')}/omnigenome_wrapper.py"
@@ -104,7 +104,9 @@ class OmniTokenizer:
             warnings.warn(
                 f"No tokenizer wrapper found in {wrapper_path} -> Exception: {e}"
             )
-            kwargs.pop("num_labels", None) # Remove num_labels if it exists, as it may not be applicable
+            kwargs.pop(
+                "num_labels", None
+            )  # Remove num_labels if it exists, as it may not be applicable
 
             tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, **kwargs)
 

omnigenome/src/dataset/omni_dataset.py

@@ -27,11 +27,11 @@ from ... import __name__, __version__
 class OmniDatasetForTokenClassification(OmniDataset):
     """
     Dataset class specifically designed for token classification tasks in genomics.
-    
+
     This class extends `OmniDataset` to provide functionalities for preparing input sequences
     and their corresponding token-level labels. It's designed for tasks where each token
     in a sequence needs to be classified independently.
-    
+
     Attributes:
         metadata: Dictionary containing dataset metadata including library information
         label2id: Mapping from label strings to integer IDs
@@ -68,7 +68,7 @@ class OmniDatasetForTokenClassification(OmniDataset):
     def prepare_input(self, instance, **kwargs):
         """
         Prepare a single data instance for token classification.
-        
+
         This method handles both string sequences and dictionary instances
         containing sequence and label information. It tokenizes the input
         sequence and prepares token-level labels for classification.
@@ -138,11 +138,11 @@ class OmniDatasetForTokenClassification(OmniDataset):
 class OmniDatasetForSequenceClassification(OmniDataset):
     """
     Dataset class for sequence classification tasks in genomics.
-    
+
     This class extends `OmniDataset` to prepare input sequences and their corresponding
     sequence-level labels. It's designed for tasks where the entire sequence needs
     to be classified into one of several categories.
-    
+
     Attributes:
         metadata: Dictionary containing dataset metadata including library information
         label2id: Mapping from label strings to integer IDs
@@ -179,7 +179,7 @@ class OmniDatasetForSequenceClassification(OmniDataset):
     def prepare_input(self, instance, **kwargs):
         """
         Prepare a single data instance for sequence classification.
-        
+
         This method handles both string sequences and dictionary instances
         containing sequence and label information. It tokenizes the input
         sequence and prepares sequence-level labels for classification.
@@ -238,11 +238,11 @@ class OmniDatasetForSequenceClassification(OmniDataset):
 class OmniDatasetForTokenRegression(OmniDataset):
     """
     Dataset class for token regression tasks in genomics.
-    
+
     This class extends `OmniDataset` to prepare input sequences and their corresponding
     token-level regression targets. It's designed for tasks where each token in a
     sequence needs to be assigned a continuous value.
-    
+
     Attributes:
         metadata: Dictionary containing dataset metadata including library information
     """
@@ -278,7 +278,7 @@ class OmniDatasetForTokenRegression(OmniDataset):
     def prepare_input(self, instance, **kwargs):
         """
         Prepare a single data instance for token regression.
-        
+
         This method handles both string sequences and dictionary instances
         containing sequence and regression target information. It tokenizes
         the input sequence and prepares token-level regression targets.
@@ -330,7 +330,9 @@ class OmniDatasetForTokenRegression(OmniDataset):
             # Handle token-level regression labels
             if isinstance(labels, (list, tuple)):
                 # Ensure labels match sequence length
-                labels = list(labels)[:self.max_length - 2]  # Account for special tokens
+                labels = list(labels)[
+                    : self.max_length - 2
+                ]  # Account for special tokens
                 labels = [-100] + labels + [-100]  # Add padding for special tokens
             else:
                 # Single value for the entire sequence
@@ -343,11 +345,11 @@ class OmniDatasetForTokenRegression(OmniDataset):
 class OmniDatasetForSequenceRegression(OmniDataset):
     """
     Dataset class for sequence regression tasks in genomics.
-    
+
     This class extends `OmniDataset` to prepare input sequences and their corresponding
     sequence-level regression targets. It's designed for tasks where the entire
     sequence needs to be assigned a continuous value.
-    
+
     Attributes:
         metadata: Dictionary containing dataset metadata including library information
     """
@@ -383,7 +385,7 @@ class OmniDatasetForSequenceRegression(OmniDataset):
     def prepare_input(self, instance, **kwargs):
         """
         Prepare a single data instance for sequence regression.
-        
+
         This method handles both string sequences and dictionary instances
         containing sequence and regression target information. It tokenizes
         the input sequence and prepares sequence-level regression targets.
@@ -432,4 +434,4 @@ class OmniDatasetForSequenceRegression(OmniDataset):
             labels = float(labels)
 
         tokenized_inputs["labels"] = torch.tensor(labels, dtype=torch.float32)
-        return tokenized_inputs
+        return tokenized_inputs

omnigenome/src/lora/__init__.py

@@ -10,4 +10,3 @@
 """
 This package contains modules for LoRA (Low-Rank Adaptation) fine-tuning.
 """
-

omnigenome/src/lora/lora_model.py

@@ -18,22 +18,23 @@ import torch
 from torch import nn
 from omnigenome.src.misc.utils import fprint
 
+
 def find_linear_target_modules(model, keyword_filter=None, use_full_path=True):
     """
     Find linear modules in a model that can be targeted for LoRA adaptation.
-    
+
     This function searches through a model's modules to identify linear layers
     that can be adapted using LoRA. It supports filtering by keyword patterns
     to target specific types of layers.
-    
+
     Args:
         model: The model to search for linear modules
         keyword_filter (str, list, tuple, optional): Keywords to filter modules by name
         use_full_path (bool): Whether to return full module paths or just names (default: True)
-        
+
     Returns:
         list: Sorted list of linear module names that can be targeted for LoRA
-        
+
     Raises:
         TypeError: If keyword_filter is not None, str, or a list/tuple of str
     """
@@ -46,31 +47,32 @@ def find_linear_target_modules(model, keyword_filter=None, use_full_path=True):
         elif not isinstance(keyword_filter, (list, tuple)):
             raise TypeError("keyword_filter must be None, str, or a list/tuple of str")
 
-        pattern = '|'.join(map(re.escape, keyword_filter))
+        pattern = "|".join(map(re.escape, keyword_filter))
 
     linear_modules = set()
     for name, module in model.named_modules():
         if isinstance(module, nn.Linear):
             if keyword_filter is None or re.search(pattern, name, re.IGNORECASE):
-                linear_modules.add(name if use_full_path else name.split('.')[-1])
+                linear_modules.add(name if use_full_path else name.split(".")[-1])
 
     return sorted(linear_modules)
 
+
 def auto_lora_model(model, **kwargs):
     """
     Automatically create a LoRA-adapted model.
-    
+
     This function automatically identifies suitable target modules and creates
     a LoRA-adapted version of the input model. It handles configuration
     setup and parameter freezing for efficient fine-tuning.
-    
+
     Args:
         model: The base model to adapt with LoRA
         **kwargs: Additional LoRA configuration parameters
-        
+
     Returns:
         The LoRA-adapted model
-        
+
     Raises:
         AssertionError: If no target modules are found for LoRA injection
     """
@@ -79,8 +81,8 @@ def auto_lora_model(model, **kwargs):
 
     # A bad case for the EVO-1 model, which has a custom config class
     ######################
-    if hasattr(model, 'config') and not isinstance(model.config, PretrainedConfig):
-        delattr(model.config, 'Loader')
+    if hasattr(model, "config") and not isinstance(model.config, PretrainedConfig):
+        delattr(model.config, "Loader")
         model.config = PretrainedConfig.from_dict(dict(model.config))
     #######################
 
@@ -92,7 +94,9 @@ def auto_lora_model(model, **kwargs):
     lora_dropout = kwargs.pop("lora_dropout", 0.1)
 
     if target_modules is None:
-        target_modules = find_linear_target_modules(model, keyword_filter=kwargs.get("keyword_filter", None))
+        target_modules = find_linear_target_modules(
+            model, keyword_filter=kwargs.get("keyword_filter", None)
+        )
     assert target_modules is not None, "No target modules found for LoRA injection."
     config = LoraConfig(
         target_modules=target_modules,
@@ -115,29 +119,30 @@ def auto_lora_model(model, **kwargs):
     )
     return lora_model
 
+
 class OmniLoraModel(nn.Module):
     """
     LoRA-adapted model for OmniGenome.
-    
+
     This class provides a wrapper around LoRA-adapted models, enabling
     efficient fine-tuning of large genomic language models while maintaining
     compatibility with the OmniGenome framework.
-    
+
     Attributes:
         lora_model: The underlying LoRA-adapted model
         config: Model configuration
         device: Device the model is running on
         dtype: Data type of the model parameters
     """
-    
+
     def __init__(self, model, **kwargs):
         """
         Initialize the LoRA-adapted model.
-        
+
         Args:
             model: The base model to adapt with LoRA
             **kwargs: LoRA configuration parameters
-            
+
         Raises:
             ValueError: If no target modules are specified for LoRA injection
         """
@@ -147,7 +152,8 @@ class OmniLoraModel(nn.Module):
             raise ValueError(
                 "No target modules found for LoRA injection. To perform LoRA adaptation fine-tuning, "
                 "please specify the target modules using the 'target_modules' argument. "
-                "The target modules depend on the model architecture, such as 'query', 'value', etc. ")
+                "The target modules depend on the model architecture, such as 'query', 'value', etc. "
+            )
 
         self.lora_model = auto_lora_model(model, **kwargs)
 
@@ -159,23 +165,23 @@ class OmniLoraModel(nn.Module):
         )
 
         self.config = model.config
-        self.to('cpu')  # Move the model to CPU initially
+        self.to("cpu")  # Move the model to CPU initially
         fprint(
             "LoRA model initialized with the following configuration:\n",
-            self.lora_model
+            self.lora_model,
         )
 
     def to(self, *args, **kwargs):
         """
         Move the model to a specific device and data type.
-        
+
         This method overrides the default to() method to ensure the LoRA model
         and its components are properly moved to the target device and dtype.
-        
+
         Args:
             *args: Device specification (e.g., 'cuda', 'cpu')
             **kwargs: Additional arguments including dtype
-            
+
         Returns:
             self: The model instance
         """
@@ -188,20 +194,20 @@ class OmniLoraModel(nn.Module):
                 break
             for module in self.lora_model.modules():
                 module.device = self.device
-                if hasattr(module, 'dtype'):
+                if hasattr(module, "dtype"):
                     module.dtype = self.dtype
         except Exception as e:
-            pass # Ignore errors if parameters are not available
+            pass  # Ignore errors if parameters are not available
         return self
 
     def forward(self, *args, **kwargs):
         """
         Forward pass through the LoRA model.
-        
+
         Args:
             *args: Positional arguments for the forward pass
             **kwargs: Keyword arguments for the forward pass
-            
+
         Returns:
             The output from the LoRA model
         """
@@ -210,11 +216,11 @@ class OmniLoraModel(nn.Module):
     def predict(self, *args, **kwargs):
         """
         Generate predictions using the LoRA model.
-        
+
         Args:
             *args: Positional arguments for prediction
             **kwargs: Keyword arguments for prediction
-            
+
         Returns:
             Model predictions
         """
@@ -223,11 +229,11 @@ class OmniLoraModel(nn.Module):
     def save(self, *args, **kwargs):
         """
         Save the LoRA model.
-        
+
         Args:
             *args: Positional arguments for saving
             **kwargs: Keyword arguments for saving
-            
+
         Returns:
             Result of the save operation
         """
@@ -236,7 +242,7 @@ class OmniLoraModel(nn.Module):
     def model_info(self):
         """
         Get information about the LoRA model.
-        
+
         Returns:
             Model information from the base model
         """
@@ -245,10 +251,10 @@ class OmniLoraModel(nn.Module):
     def set_loss_fn(self, fn):
         """
         Set the loss function for the LoRA model.
-        
+
         Args:
             fn: Loss function to set
-            
+
         Returns:
             Result of setting the loss function
         """
@@ -257,10 +263,10 @@ class OmniLoraModel(nn.Module):
     def last_hidden_state_forward(self, **kwargs):
         """
         Forward pass to get the last hidden state.
-        
+
         Args:
             **kwargs: Keyword arguments for the forward pass
-            
+
         Returns:
             Last hidden state from the base model
         """
@@ -269,7 +275,7 @@ class OmniLoraModel(nn.Module):
     def tokenizer(self):
         """
         Get the tokenizer from the base model.
-        
+
         Returns:
             The tokenizer from the base model
         """
@@ -278,7 +284,7 @@ class OmniLoraModel(nn.Module):
     def config(self):
         """
         Get the configuration from the base model.
-        
+
         Returns:
             The configuration from the base model
         """
@@ -287,8 +293,8 @@ class OmniLoraModel(nn.Module):
     def model(self):
         """
         Get the base model.
-        
+
         Returns:
             The base model
         """
-        return self.lora_model.base_model.model
+        return self.lora_model.base_model.model

omnigenome/src/metric/classification_metric.py

@@ -19,17 +19,17 @@ from ..abc.abstract_metric import OmniMetric
 class ClassificationMetric(OmniMetric):
     """
     Classification metric class for evaluating classification models.
-    
+
     This class provides a comprehensive interface for classification metrics
     in the OmniGenome framework. It integrates with scikit-learn's classification
     metrics and provides additional functionality for handling genomic classification
     tasks.
-    
+
     The class automatically exposes all scikit-learn classification metrics as
     callable attributes, making them easily accessible for evaluation. It also
     handles special cases like Hugging Face's EvalPrediction objects and
     provides proper handling of ignored labels.
-    
+
     Attributes:
         metric_func (callable): A callable metric function from sklearn.metrics.
         ignore_y (any): A value in the ground truth labels to be ignored during
@@ -42,10 +42,10 @@ class ClassificationMetric(OmniMetric):
         Initializes the classification metric.
 
         Args:
-            metric_func (callable, optional): A callable metric function from 
+            metric_func (callable, optional): A callable metric function from
                                             sklearn.metrics. If None, subclasses
                                             should implement their own compute method.
-            ignore_y (any, optional): A value in the ground truth labels to be 
+            ignore_y (any, optional): A value in the ground truth labels to be
                                     ignored during metric computation. Defaults to -100.
             *args: Additional positional arguments.
             **kwargs: Additional keyword arguments.
@@ -53,7 +53,7 @@ class ClassificationMetric(OmniMetric):
         Example:
             >>> # Initialize with a specific metric function
             >>> metric = ClassificationMetric(metrics.accuracy_score)
-            
+
             >>> # Initialize with ignore value
             >>> metric = ClassificationMetric(ignore_y=-100)
         """
@@ -64,7 +64,7 @@ class ClassificationMetric(OmniMetric):
     def __getattribute__(self, name):
         """
         Custom attribute getter that provides dynamic access to scikit-learn metrics.
-        
+
         This method provides transparent access to all scikit-learn classification
         metrics. When a metric function is accessed, it returns a callable wrapper
         that handles the metric computation with proper preprocessing.
@@ -91,7 +91,7 @@ class ClassificationMetric(OmniMetric):
             def wrapper(y_true=None, y_pred=None, *args, **kwargs):
                 """
                 Compute the metric, based on the true and predicted values.
-                
+
                 This wrapper function handles various input formats including
                 Hugging Face's EvalPrediction objects and provides proper
                 preprocessing for metric computation.
@@ -99,7 +99,7 @@ class ClassificationMetric(OmniMetric):
                 Args:
                     y_true: The true values (ground truth labels).
                     y_pred: The predicted values (model predictions).
-                    ignore_y: The value to ignore in the predictions and true 
+                    ignore_y: The value to ignore in the predictions and true
                              values in corresponding positions.
                     *args: Additional positional arguments for the metric function.
                     **kwargs: Additional keyword arguments for the metric function.
@@ -111,7 +111,7 @@ class ClassificationMetric(OmniMetric):
                     >>> # Standard usage
                     >>> result = accuracy_fn(y_true, y_pred)
                     >>> print(result)  # {'accuracy_score': 0.85}
-                    
+
                     >>> # With Hugging Face EvalPrediction
                     >>> result = accuracy_fn(eval_prediction)
                     >>> print(result)  # {'accuracy_score': 0.85}
@@ -152,7 +152,7 @@ class ClassificationMetric(OmniMetric):
     def compute(self, y_true, y_pred, *args, **kwargs):
         """
         Compute the metric, based on the true and predicted values.
-        
+
         This method computes the classification metric using the provided
         metric function. It handles preprocessing and applies any additional
         keyword arguments.