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

omnigenome/src/tokenizer/kmers_tokenizer.py

@@ -16,18 +16,18 @@ warnings.filterwarnings("once")
 class OmniKmersTokenizer(OmniTokenizer):
     """
     A k-mer based tokenizer for genomic sequences.
-    
+
     This tokenizer breaks genomic sequences into overlapping k-mers and uses
     a base tokenizer to convert them into token IDs. It supports various
     k-mer sizes and overlap configurations for different genomic applications.
-    
+
     Attributes:
         base_tokenizer: The underlying tokenizer for converting k-mers to IDs
         k: Size of k-mers
         overlap: Number of overlapping positions between consecutive k-mers
         max_length: Maximum sequence length for tokenization
         metadata: Dictionary containing tokenizer metadata
-        
+
     Example:
         >>> from omnigenome.src.tokenizer import OmniKmersTokenizer
         >>> from transformers import AutoTokenizer
@@ -42,7 +42,7 @@ class OmniKmersTokenizer(OmniTokenizer):
     def __init__(self, base_tokenizer=None, k=3, overlap=0, max_length=512, **kwargs):
         """
         Initialize the OmniKmersTokenizer.
-        
+
         Args:
             base_tokenizer: The base tokenizer for converting k-mers to token IDs
             k (int, optional): Size of k-mers. Defaults to 3
@@ -59,18 +59,18 @@ class OmniKmersTokenizer(OmniTokenizer):
     def __call__(self, sequence, **kwargs):
         """
         Tokenize a sequence or list of sequences into tokenized inputs.
-        
+
         This method processes the input sequence(s) by first converting them to k-mers,
         then using the base tokenizer to convert k-mers to token IDs. It handles
         sequence preprocessing (U/T conversion) and adds special tokens.
-        
+
         Args:
             sequence (str or list): Input sequence(s) to tokenize
             **kwargs: Additional keyword arguments including max_length
-            
+
         Returns:
             dict: Dictionary containing tokenized inputs with keys 'input_ids' and 'attention_mask'
-            
+
         Example:
             >>> sequence = "ACGUAGGUAUCGUAGA"
             >>> tokenized = tokenizer(sequence)
@@ -126,14 +126,14 @@ class OmniKmersTokenizer(OmniTokenizer):
     def from_pretrained(model_name_or_path, **kwargs):
         """
         Create a k-mers tokenizer from a pre-trained model.
-        
+
         Args:
             model_name_or_path (str): Name or path of the pre-trained model
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             OmniKmersTokenizer: Initialized k-mers tokenizer
-            
+
         Example:
             >>> tokenizer = OmniKmersTokenizer.from_pretrained("facebook/esm2_t12_35M_UR50D")
             >>> print(type(tokenizer))
@@ -149,17 +149,17 @@ class OmniKmersTokenizer(OmniTokenizer):
     def tokenize(self, sequence, **kwargs):
         """
         Convert sequence(s) into k-mers.
-        
+
         This method breaks the input sequence(s) into overlapping k-mers based on
         the configured k-mer size and overlap parameters.
-        
+
         Args:
             sequence (str or list): Input sequence(s) to convert to k-mers
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             list: List of k-mer lists for each input sequence
-            
+
         Example:
             >>> sequence = "ACGUAGGUAUCGUAGA"
             >>> k_mers = tokenizer.tokenize(sequence)
@@ -184,11 +184,11 @@ class OmniKmersTokenizer(OmniTokenizer):
     def encode(self, input_ids, **kwargs):
         """
         Encode input IDs using the base tokenizer.
-        
+
         Args:
             input_ids: Input IDs to encode
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             Encoded input IDs
         """
@@ -197,11 +197,11 @@ class OmniKmersTokenizer(OmniTokenizer):
     def decode(self, input_ids, **kwargs):
         """
         Decode input IDs using the base tokenizer.
-        
+
         Args:
             input_ids: Input IDs to decode
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             Decoded sequence
         """
@@ -210,13 +210,13 @@ class OmniKmersTokenizer(OmniTokenizer):
     def encode_plus(self, sequence, **kwargs):
         """
         Encode a sequence with additional information.
-        
+
         This method is not yet implemented for k-mers tokenizer.
-        
+
         Args:
             sequence: Input sequence
             **kwargs: Additional keyword arguments
-            
+
         Raises:
             NotImplementedError: This method is not implemented yet
         """

omnigenome/src/tokenizer/single_nucleotide_tokenizer.py

@@ -19,16 +19,16 @@ warnings.filterwarnings("once")
 class OmniSingleNucleotideTokenizer(OmniTokenizer):
     """
     Tokenizer for single nucleotide tokenization in genomics.
-    
+
     This tokenizer converts genomic sequences into individual nucleotide tokens,
     where each nucleotide (A, T, C, G, U) becomes a separate token. It's designed
     for genomic sequence processing where fine-grained nucleotide-level analysis
     is required.
-    
+
     The tokenizer supports various preprocessing options including U/T conversion
     and whitespace addition between nucleotides. It also handles special tokens
     like BOS (beginning of sequence) and EOS (end of sequence) tokens.
-    
+
     Attributes:
         u2t (bool): Whether to convert 'U' to 'T'.
         t2u (bool): Whether to convert 'T' to 'U'.
@@ -54,7 +54,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def __call__(self, sequence, **kwargs):
         """
         Tokenizes sequences using single nucleotide tokenization.
-        
+
         This method converts genomic sequences into tokenized inputs suitable
         for model training and inference. It handles sequence preprocessing,
         tokenization, and padding/truncation.
@@ -76,7 +76,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
             >>> # Tokenize a single sequence
             >>> inputs = tokenizer("ATCGATCG")
             >>> print(inputs['input_ids'].shape)  # torch.Size([1, seq_len])
-            
+
             >>> # Tokenize multiple sequences
             >>> inputs = tokenizer(["ATCGATCG", "GCTAGCTA"])
             >>> print(inputs['input_ids'].shape)  # torch.Size([2, seq_len])
@@ -134,7 +134,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def from_pretrained(model_name_or_path, **kwargs):
         """
         Loads a single nucleotide tokenizer from a pre-trained model.
-        
+
         This method creates a single nucleotide tokenizer wrapper around
         a Hugging Face tokenizer loaded from a pre-trained model.
 
@@ -156,7 +156,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def tokenize(self, sequence, **kwargs):
         """
         Converts a sequence into a list of individual nucleotide tokens.
-        
+
         This method tokenizes genomic sequences by treating each nucleotide
         as a separate token. It handles both single sequences and lists of sequences.
 
@@ -172,7 +172,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
             >>> # Tokenize a single sequence
             >>> tokens = tokenizer.tokenize("ATCGATCG")
             >>> print(tokens)  # [['A', 'T', 'C', 'G', 'A', 'T', 'C', 'G']]
-            
+
             >>> # Tokenize multiple sequences
             >>> tokens = tokenizer.tokenize(["ATCGATCG", "GCTAGCTA"])
             >>> print(tokens)  # [['A', 'T', 'C', 'G', ...], ['G', 'C', 'T', 'A', ...]]
@@ -191,7 +191,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def encode(self, sequence, **kwargs):
         """
         Converts a sequence into a list of token IDs.
-        
+
         This method encodes genomic sequences into token IDs using the
         underlying base tokenizer.
 
@@ -211,7 +211,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def decode(self, sequence, **kwargs):
         """
         Converts a list of token IDs back into a sequence.
-        
+
         This method decodes token IDs back into genomic sequences using
         the underlying base tokenizer.
 
@@ -231,7 +231,7 @@ class OmniSingleNucleotideTokenizer(OmniTokenizer):
     def encode_plus(self, sequence, **kwargs):
         """
         Encodes a sequence with additional information.
-        
+
         This method provides enhanced encoding with additional information
         like attention masks and token type IDs.
 

omnigenome/src/trainer/accelerate_trainer.py

@@ -21,15 +21,15 @@ from ..misc.utils import env_meta_info, fprint, seed_everything
 def _infer_optimization_direction(metrics, prev_metrics):
     """
     Infer the optimization direction based on metric values.
-    
+
     This function analyzes the trend of metric values to determine whether
     larger values are better (e.g., accuracy) or smaller values are better
     (e.g., loss).
-    
+
     Args:
         metrics (dict): Current metric values
         prev_metrics (list): Previous metric values
-        
+
     Returns:
         str: Either 'larger_is_better' or 'smaller_is_better'
     """
@@ -91,11 +91,11 @@ def _infer_optimization_direction(metrics, prev_metrics):
 class AccelerateTrainer:
     """
     A distributed training trainer using HuggingFace Accelerate.
-    
+
     This trainer provides distributed training capabilities with automatic mixed precision,
     gradient accumulation, and early stopping. It supports both single and multi-GPU
     training with seamless integration with HuggingFace Accelerate.
-    
+
     Attributes:
         model: The model to train
         train_loader: DataLoader for training data
@@ -110,7 +110,7 @@ class AccelerateTrainer:
         accelerator: HuggingFace Accelerate instance
         metrics: Dictionary to store training metrics
         predictions: Dictionary to store predictions
-        
+
     Example:
         >>> from omnigenome.src.trainer import AccelerateTrainer
         >>> trainer = AccelerateTrainer(
@@ -143,7 +143,7 @@ class AccelerateTrainer:
     ):
         """
         Initialize the AccelerateTrainer.
-        
+
         Args:
             model: The model to train
             train_dataset (torch.utils.data.Dataset, optional): Training dataset
@@ -293,14 +293,14 @@ class AccelerateTrainer:
     def evaluate(self):
         """
         Evaluate the model on the validation dataset.
-        
+
         This method runs the model in evaluation mode and computes metrics
         on the validation dataset. It handles distributed evaluation and
         gathers results from all processes.
-        
+
         Returns:
             dict: Dictionary containing evaluation metrics
-            
+
         Example:
             >>> metrics = trainer.evaluate()
             >>> print(f"Validation accuracy: {metrics['accuracy']:.4f}")
@@ -364,14 +364,14 @@ class AccelerateTrainer:
     def test(self):
         """
         Test the model on the test dataset.
-        
+
         This method runs the model in evaluation mode and computes metrics
         on the test dataset. It handles distributed testing and gathers
         results from all processes.
-        
+
         Returns:
             dict: Dictionary containing test metrics
-            
+
         Example:
             >>> metrics = trainer.test()
             >>> print(f"Test accuracy: {metrics['accuracy']:.4f}")
@@ -431,18 +431,18 @@ class AccelerateTrainer:
     def train(self, path_to_save=None, **kwargs):
         """
         Train the model using distributed training.
-        
+
         This method performs the complete training loop with validation,
         early stopping, and model checkpointing. It handles distributed
         training across multiple GPUs and processes.
-        
+
         Args:
             path_to_save (str, optional): Path to save the trained model
             **kwargs: Additional keyword arguments for model saving
-            
+
         Returns:
             dict: Dictionary containing training metrics
-            
+
         Example:
             >>> metrics = trainer.train(path_to_save="./checkpoints/model")
             >>> print(f"Best validation accuracy: {metrics['best_valid']['accuracy']:.4f}")
@@ -489,12 +489,20 @@ class AccelerateTrainer:
                 if "loss" not in outputs:
                     # Generally, the model should return a loss in the outputs via OmniGenBench
                     # For the Lora models, the loss is computed separately
-                    if hasattr(self.model, "loss_function") and callable(self.model.loss_function):
-                        loss = self.model.loss_function(outputs['logits'], outputs["labels"])
-                    elif (hasattr(self.model, "model")
-                          and hasattr(self.model.model, "loss_function")
-                          and callable(self.model.model.loss_function)):
-                        loss = self.model.model.loss_function(outputs['logits'], outputs["labels"])
+                    if hasattr(self.model, "loss_function") and callable(
+                        self.model.loss_function
+                    ):
+                        loss = self.model.loss_function(
+                            outputs["logits"], outputs["labels"]
+                        )
+                    elif (
+                        hasattr(self.model, "model")
+                        and hasattr(self.model.model, "loss_function")
+                        and callable(self.model.model.loss_function)
+                    ):
+                        loss = self.model.model.loss_function(
+                            outputs["logits"], outputs["labels"]
+                        )
                     else:
                         raise ValueError(
                             "The model does not have a loss function defined. "
@@ -585,11 +593,11 @@ class AccelerateTrainer:
     def _is_metric_better(self, metrics, stage="valid"):
         """
         Check if the current metrics are better than the best metrics so far.
-        
+
         Args:
             metrics (dict): Current metrics
             stage (str): Stage of evaluation ('valid' or 'test')
-            
+
         Returns:
             bool: True if current metrics are better, False otherwise
         """
@@ -643,10 +651,10 @@ class AccelerateTrainer:
     def predict(self, data_loader):
         """
         Make predictions using the trained model.
-        
+
         Args:
             data_loader: DataLoader containing data to predict on
-            
+
         Returns:
             dict: Dictionary containing predictions
         """
@@ -655,10 +663,10 @@ class AccelerateTrainer:
     def get_model(self, **kwargs):
         """
         Get the trained model.
-        
+
         Args:
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             The trained model
         """
@@ -667,10 +675,10 @@ class AccelerateTrainer:
     def compute_metrics(self):
         """
         Compute metrics for evaluation.
-        
+
         This method should be implemented by subclasses to provide specific
         metric computation logic.
-        
+
         Raises:
             NotImplementedError: If compute_metrics method is not implemented
         """
@@ -682,7 +690,7 @@ class AccelerateTrainer:
     def save_model(self, path, overwrite=False, **kwargs):
         """
         Save the trained model.
-        
+
         Args:
             path (str): Path to save the model
             overwrite (bool, optional): Whether to overwrite existing files. Defaults to False

omnigenome/src/trainer/hf_trainer.py

@@ -24,19 +24,19 @@ from ... import __version__ as omnigenome_version
 class HFTrainer(Trainer):
     """
     HuggingFace trainer wrapper for OmniGenome models.
-    
+
     This class extends the HuggingFace Trainer to include OmniGenome-specific
     metadata and functionality while maintaining full compatibility with the
     HuggingFace training ecosystem.
-    
+
     Attributes:
         metadata: Dictionary containing OmniGenome library information
     """
-    
+
     def __init__(self, *args, **kwargs):
         """
         Initialize the HuggingFace trainer wrapper.
-        
+
         Args:
             *args: Positional arguments passed to the parent Trainer
             **kwargs: Keyword arguments passed to the parent Trainer
@@ -51,19 +51,19 @@ class HFTrainer(Trainer):
 class HFTrainingArguments(TrainingArguments):
     """
     HuggingFace training arguments wrapper for OmniGenome models.
-    
+
     This class extends the HuggingFace TrainingArguments to include
     OmniGenome-specific metadata while maintaining full compatibility
     with the HuggingFace training ecosystem.
-    
+
     Attributes:
         metadata: Dictionary containing OmniGenome library information
     """
-    
+
     def __init__(self, *args, **kwargs):
         """
         Initialize the HuggingFace training arguments wrapper.
-        
+
         Args:
             *args: Positional arguments passed to the parent TrainingArguments
             **kwargs: Keyword arguments passed to the parent TrainingArguments

omnigenome/src/trainer/trainer.py

@@ -29,14 +29,14 @@ from torch.cuda.amp import GradScaler
 def _infer_optimization_direction(metrics, prev_metrics):
     """
     Infer the optimization direction based on metric names and trends.
-    
+
     This function determines whether larger or smaller values are better for
     the given metrics by analyzing metric names and their trends over time.
-    
+
     Args:
         metrics (dict): Current metric values
         prev_metrics (list): Previous metric values from multiple epochs
-        
+
     Returns:
         str: Either "larger_is_better" or "smaller_is_better"
     """
@@ -98,11 +98,11 @@ def _infer_optimization_direction(metrics, prev_metrics):
 class Trainer:
     """
     Comprehensive trainer for OmniGenome models.
-    
+
     This trainer provides a complete training framework with automatic mixed precision,
     early stopping, metric tracking, and model checkpointing. It supports various
     training configurations and can handle different types of genomic sequence tasks.
-    
+
     Attributes:
         model: The model to be trained
         train_loader: DataLoader for training data
@@ -118,7 +118,7 @@ class Trainer:
         metrics: Dictionary to store training metrics
         predictions: Dictionary to store model predictions
     """
-    
+
     def __init__(
         self,
         model,
@@ -139,7 +139,7 @@ class Trainer:
     ):
         """
         Initialize the trainer.
-        
+
         Args:
             model: The model to be trained
             train_dataset: Training dataset
@@ -191,7 +191,9 @@ class Trainer:
         )
         self.seed = seed
         self.device = device if device else autocuda.auto_cuda()
-        self.device = torch.device(self.device) if isinstance(self.device, str) else self.device
+        self.device = (
+            torch.device(self.device) if isinstance(self.device, str) else self.device
+        )
 
         self.fast_dtype = {
             "float32": torch.float32,
@@ -218,11 +220,11 @@ class Trainer:
     def _is_metric_better(self, metrics, stage="valid"):
         """
         Check if the current metrics are better than the best metrics so far.
-        
+
         Args:
             metrics (dict): Current metric values
             stage (str): Stage name ("valid" or "test")
-            
+
         Returns:
             bool: True if current metrics are better than best metrics
         """
@@ -268,11 +270,11 @@ class Trainer:
     def train(self, path_to_save=None, **kwargs):
         """
         Train the model.
-        
+
         Args:
             path_to_save (str, optional): Path to save the best model
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             dict: Training metrics and results
         """
@@ -300,19 +302,29 @@ class Trainer:
                     self.optimizer.zero_grad()
 
                 if self.fast_dtype:
-                    with torch.autocast(device_type=self.device.type, dtype=self.fast_dtype):
+                    with torch.autocast(
+                        device_type=self.device.type, dtype=self.fast_dtype
+                    ):
                         outputs = self.model(**batch)
                 else:
                     outputs = self.model(**batch)
                 if "loss" not in outputs:
                     # Generally, the model should return a loss in the outputs via OmniGenBench
                     # For the Lora models, the loss is computed separately
-                    if hasattr(self.model, "loss_function") and callable(self.model.loss_function):
-                        loss = self.model.loss_function(outputs['logits'], outputs["labels"])
-                    elif (hasattr(self.model, "model")
-                          and hasattr(self.model.model, "loss_function")
-                          and callable(self.model.model.loss_function)):
-                        loss = self.model.model.loss_function(outputs['logits'], outputs["labels"])
+                    if hasattr(self.model, "loss_function") and callable(
+                        self.model.loss_function
+                    ):
+                        loss = self.model.loss_function(
+                            outputs["logits"], outputs["labels"]
+                        )
+                    elif (
+                        hasattr(self.model, "model")
+                        and hasattr(self.model.model, "loss_function")
+                        and callable(self.model.model.loss_function)
+                    ):
+                        loss = self.model.model.loss_function(
+                            outputs["logits"], outputs["labels"]
+                        )
                     else:
                         raise ValueError(
                             "The model does not have a loss function defined. "
@@ -480,10 +492,10 @@ class Trainer:
     def get_model(self, **kwargs):
         """
         Get the trained model.
-        
+
         Args:
             **kwargs: Additional keyword arguments
-            
+
         Returns:
             The trained model
         """
@@ -492,7 +504,7 @@ class Trainer:
     def compute_metrics(self):
         """
         Get the metric computation functions.
-        
+
         Returns:
             list: List of metric computation functions
         """
@@ -501,10 +513,10 @@ class Trainer:
     def unwrap_model(self, model=None):
         """
         Unwrap the model from any distributed training wrappers.
-        
+
         Args:
             model: Model to unwrap (default: None, uses self.model)
-            
+
         Returns:
             The unwrapped model
         """
@@ -538,7 +550,7 @@ class Trainer:
         """
         if os.path.exists(self._model_state_dict_path):
             self.unwrap_model().load_state_dict(
-                torch.load(self._model_state_dict_path, map_location='cpu')
+                torch.load(self._model_state_dict_path, map_location="cpu")
             )
             self.unwrap_model().to(self.device)
 

omnigenome/utility/dataset_hub/__init__.py

@@ -10,4 +10,3 @@
 """
 This package contains modules for the dataset hub.
 """
-