omnigenome 0.3.0a0__py3-none-any.whl

omnigenome/src/lora/lora_model.py

@@ -0,0 +1,294 @@
+# -*- coding: utf-8 -*-
+# file: lora_model.py
+# time: 12:36 11/06/2025
+# author: YANG, HENG <hy345@exeter.ac.uk> (杨恒)
+# homepage: https://yangheng95.github.io
+# 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-2025. All Rights Reserved.
+"""
+Low-Rank Adaptation (LoRA) models for OmniGenome.
+
+This module provides LoRA implementation for efficient fine-tuning of large
+genomic language models. LoRA reduces the number of trainable parameters
+by adding low-rank adaptation layers to existing model weights.
+"""
+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
+    """
+    import re
+    from torch import nn
+
+    if keyword_filter is not None:
+        if isinstance(keyword_filter, str):
+            keyword_filter = [keyword_filter]
+        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))
+
+    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])
+
+    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
+    """
+    from peft import LoraConfig, get_peft_model
+    from transformers import PretrainedConfig
+
+    # 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')
+        model.config = PretrainedConfig.from_dict(dict(model.config))
+    #######################
+
+    target_modules = kwargs.pop("target_modules", None)
+    use_rslora = kwargs.pop("use_rslora", True)
+    bias = kwargs.pop("bias", "none")
+    r = kwargs.pop("r", 32)
+    lora_alpha = kwargs.pop("lora_alpha", 256)
+    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))
+    assert target_modules is not None, "No target modules found for LoRA injection."
+    config = LoraConfig(
+        target_modules=target_modules,
+        r=r,
+        lora_alpha=lora_alpha,
+        lora_dropout=lora_dropout,
+        bias=bias,
+        use_rslora=use_rslora,
+        **kwargs,
+    )
+
+    for param in model.parameters():
+        param.requires_grad = False
+
+    lora_model = get_peft_model(model, config)
+    trainable_params, all_param = lora_model.get_nb_trainable_parameters()
+    fprint(
+        f"trainable params: {trainable_params:,d} || all params: {all_param:,d}"
+        f" || trainable%: {100 * trainable_params / all_param:.4f}"
+    )
+    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
+        """
+        super(OmniLoraModel, self).__init__()
+        target_modules = kwargs.get("target_modules", None)
+        if target_modules is None:
+            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. ")
+
+        self.lora_model = auto_lora_model(model, **kwargs)
+
+        fprint(
+            "To reduce GPU memory occupation, "
+            "you should avoid include non-trainable parameters into optimizers, "
+            "e.g.,  optimizer = torch.optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), ...), "
+            "AVOIDING: optimizer = torch.optim.AdamW(model.parameters(), ...)"
+        )
+
+        self.config = model.config
+        self.to('cpu')  # Move the model to CPU initially
+        fprint(
+            "LoRA model initialized with the following configuration:\n",
+            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
+        """
+        self.lora_model.to(*args, **kwargs)
+        try:
+            # For evo-1 and similar models, we need to set the device and dtype
+            for param in self.parameters():
+                self.device = param.device
+                self.dtype = param.dtype
+                break
+            for module in self.lora_model.modules():
+                module.device = self.device
+                if hasattr(module, 'dtype'):
+                    module.dtype = self.dtype
+        except Exception as e:
+            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
+        """
+        return self.lora_model(*args, **kwargs)
+
+    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
+        """
+        return self.lora_model.base_model.predict(*args, **kwargs)
+
+    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
+        """
+        return self.lora_model.base_model.save(*args, **kwargs)
+
+    def model_info(self):
+        """
+        Get information about the LoRA model.
+        
+        Returns:
+            Model information from the base model
+        """
+        return self.lora_model.base_model.model_info()
+
+    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
+        """
+        return self.lora_model.base_model.set_loss_fn(fn)
+
+    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
+        """
+        return self.lora_model.base_model.last_hidden_state_forward(**kwargs)
+
+    def tokenizer(self):
+        """
+        Get the tokenizer from the base model.
+        
+        Returns:
+            The tokenizer from the base model
+        """
+        return self.lora_model.base_model.tokenizer
+
+    def config(self):
+        """
+        Get the configuration from the base model.
+        
+        Returns:
+            The configuration from the base model
+        """
+        return self.lora_model.base_model.config
+
+    def model(self):
+        """
+        Get the base model.
+        
+        Returns:
+            The base model
+        """
+        return self.lora_model.base_model.model

omnigenome/src/metric/__init__.py

@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+# file: __init__.py
+# time: 12:53 09/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 evaluation metrics.
+"""
+
+from .classification_metric import ClassificationMetric
+from .ranking_metric import RankingMetric
+from .regression_metric import RegressionMetric

omnigenome/src/metric/classification_metric.py

@@ -0,0 +1,184 @@
+# -*- coding: utf-8 -*-
+# file: classification_metric.py
+# time: 12:57 09/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.
+
+import types
+import warnings
+
+import numpy as np
+import sklearn.metrics as metrics
+
+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
+                       metric computation. Defaults to -100.
+        kwargs (dict): Additional keyword arguments for metric computation.
+    """
+
+    def __init__(self, metric_func=None, ignore_y=-100, *args, **kwargs):
+        """
+        Initializes the classification metric.
+
+        Args:
+            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 
+                                    ignored during metric computation. Defaults to -100.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+
+        Example:
+            >>> # Initialize with a specific metric function
+            >>> metric = ClassificationMetric(metrics.accuracy_score)
+            
+            >>> # Initialize with ignore value
+            >>> metric = ClassificationMetric(ignore_y=-100)
+        """
+        super().__init__(metric_func, ignore_y, *args, **kwargs)
+        self.kwargs = kwargs
+
+    # def __getattr__(self, name):
+    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.
+
+        Args:
+            name (str): The attribute name to get.
+
+        Returns:
+            callable: A wrapper function for the requested metric, or the original
+                     attribute if it's not a metric function.
+
+        Example:
+            >>> metric = ClassificationMetric()
+            >>> # Access any scikit-learn metric
+            >>> accuracy_fn = metric.accuracy_score
+            >>> result = accuracy_fn(y_true, y_pred)
+        """
+        # Get the metric function
+        metric_func = getattr(metrics, name, None)
+        if metric_func and isinstance(metric_func, types.FunctionType):
+            setattr(self, "compute", metric_func)
+            # If the metric function exists, return a wrapper function
+
+            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.
+
+                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 
+                             values in corresponding positions.
+                    *args: Additional positional arguments for the metric function.
+                    **kwargs: Additional keyword arguments for the metric function.
+
+                Returns:
+                    dict: A dictionary with the metric name as key and its value.
+
+                Example:
+                    >>> # 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}
+                """
+
+                # This is an ugly method to handle the case when the predictions are in the form of a tuple
+                # for huggingface trainers
+                if y_true.__class__.__name__ == "EvalPrediction":
+                    eval_prediction = y_true
+                    if hasattr(eval_prediction, "label_ids"):
+                        y_true = eval_prediction.label_ids
+                    if hasattr(eval_prediction, "labels"):
+                        y_true = eval_prediction.labels
+                    predictions = eval_prediction.predictions
+                    for i in range(len(predictions)):
+                        if predictions[i].shape == y_true.shape and not np.all(
+                            predictions[i] == y_true
+                        ):
+                            y_score = predictions[i]
+                            break
+
+                y_true, y_pred = ClassificationMetric.flatten(y_true, y_pred)
+                y_true_mask_idx = np.where(y_true != self.ignore_y)
+                if self.ignore_y is not None:
+                    y_true = y_true[y_true_mask_idx]
+                    try:
+                        y_pred = y_pred[y_true_mask_idx]
+                    except Exception as e:
+                        warnings.warn(str(e))
+
+                kwargs.update(self.kwargs)
+                return {name: self.compute(y_true, y_pred, *args, **kwargs)}
+
+            return wrapper
+        else:
+            return super().__getattribute__(name)
+
+    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.
+
+        Args:
+            y_true: The true values (ground truth labels).
+            y_pred: The predicted values (model predictions).
+            *args: Additional positional arguments for the metric function.
+            **kwargs: Additional keyword arguments for the metric function.
+
+        Returns:
+            dict: A dictionary with the metric name as key and its value.
+
+        Raises:
+            NotImplementedError: If no metric function is provided and the method
+                              is not implemented by the subclass.
+
+        Example:
+            >>> metric = ClassificationMetric(metrics.accuracy_score)
+            >>> result = metric.compute(y_true, y_pred)
+            >>> print(result)  # {'accuracy_score': 0.85}
+        """
+        if self.metric_func is not None:
+            kwargs.update(self.kwargs)
+            return self.metric_func(y_true, y_pred, *args, **kwargs)
+        else:
+            raise NotImplementedError(
+                "Method compute() is not implemented in the child class."
+            )

omnigenome/src/metric/metric.py

@@ -0,0 +1,199 @@
+# -*- coding: utf-8 -*-
+# file: regression_metric.py
+# time: 12:57 09/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.
+
+
+import types
+import warnings
+
+import numpy as np
+import sklearn.metrics as metrics
+
+from ..abc.abstract_metric import OmniMetric
+
+
+def mcrmse(y_true, y_pred):
+    """
+    Compute Mean Column Root Mean Square Error (MCRMSE).
+    
+    MCRMSE is a multi-target regression metric that computes the RMSE for each target
+    column and then takes the mean across all targets.
+    
+    Args:
+        y_true (np.ndarray): Ground truth values with shape (n_samples, n_targets)
+        y_pred (np.ndarray): Predicted values with shape (n_samples, n_targets)
+        
+    Returns:
+        float: Mean Column Root Mean Square Error
+        
+    Raises:
+        ValueError: If y_true and y_pred have different shapes
+        
+    Example:
+        >>> y_true = np.array([[1, 2], [3, 4], [5, 6]])
+        >>> y_pred = np.array([[1.1, 2.1], [2.9, 4.1], [5.2, 5.8]])
+        >>> mcrmse(y_true, y_pred)
+        0.1833...
+    """
+    if y_true.shape != y_pred.shape:
+        raise ValueError("y_true and y_pred must have the same shape")
+    mask = y_true != -100
+    filtered_y_pred = y_pred[mask]
+    filtered_y_true = y_true[mask]
+    rmse_per_target = np.sqrt(np.mean((filtered_y_true - filtered_y_pred) ** 2, axis=0))
+    mcrmse_value = np.mean(rmse_per_target)
+    return mcrmse_value
+
+
+setattr(metrics, "mcrmse", mcrmse)
+
+
+class Metric(OmniMetric):
+    """
+    A flexible metric class that provides access to all scikit-learn metrics
+    and custom metrics for evaluation.
+    
+    This class dynamically wraps scikit-learn metrics and provides a unified
+    interface for computing various evaluation metrics. It handles different
+    input formats including HuggingFace trainer outputs and supports
+    custom metric functions.
+    
+    Attributes:
+        metric_func: Custom metric function if provided
+        ignore_y: Value to ignore in predictions and true values
+        kwargs: Additional keyword arguments for metric computation
+        metrics: Dictionary of available metrics including custom ones
+        
+    Example:
+        >>> from omnigenome.src.metric import Metric
+        >>> metric = Metric(ignore_y=-100)
+        >>> y_true = [0, 1, 2, 0, 1]
+        >>> y_pred = [0, 1, 1, 0, 1]
+        >>> result = metric.accuracy(y_true, y_pred)
+        >>> print(result)
+        {'accuracy': 0.8}
+    """
+
+    def __init__(self, metric_func=None, ignore_y=-100, *args, **kwargs):
+        """
+        Initialize the Metric class.
+        
+        Args:
+            metric_func (callable, optional): Custom metric function to use
+            ignore_y (int, optional): Value to ignore in predictions and true values. Defaults to -100
+            *args: Additional positional arguments
+            **kwargs: Additional keyword arguments for metric computation
+        """
+        super().__init__(metric_func, ignore_y, *args, **kwargs)
+        self.kwargs = kwargs
+        self.metrics = {"mcrmse": mcrmse}
+        for key, value in metrics.__dict__.items():
+            setattr(self, key, value)
+
+    def __getattribute__(self, name):
+        """
+        Dynamically create metric computation methods.
+        
+        This method intercepts attribute access and creates wrapper functions
+        for scikit-learn metrics, handling different input formats and
+        preprocessing the data appropriately.
+        
+        Args:
+            name (str): Name of the metric to access
+            
+        Returns:
+            callable: Wrapper function for the requested metric
+        """
+        # Get the metric function
+        metric_func = getattr(metrics, name, None)
+
+        if metric_func and isinstance(metric_func, types.FunctionType):
+            setattr(self, "compute", metric_func)
+            # If the metric function exists, return a wrapper function
+
+            def wrapper(y_true=None, y_score=None, *args, **kwargs):
+                """
+                Compute the metric, based on the true and predicted values.
+                
+                This wrapper handles different input formats including HuggingFace
+                trainer outputs and performs necessary preprocessing.
+                
+                Args:
+                    y_true: The true values or HuggingFace EvalPrediction object
+                    y_score: The predicted values
+                    ignore_y: The value to ignore in the predictions and true values in corresponding positions
+                    *args: Additional positional arguments for the metric
+                    **kwargs: Additional keyword arguments for the metric
+                    
+                Returns:
+                    dict: Dictionary containing the metric name and computed value
+                    
+                Raises:
+                    ValueError: If neither y_true nor y_score is provided
+                """
+                # This is an ugly method to handle the case when the predictions are in the form of a tuple
+                # for huggingface trainers
+                if y_true is not None and y_score is None:
+                    if hasattr(y_true, "predictions"):
+                        y_score = y_true.predictions
+                    if hasattr(y_true, "label_ids"):
+                        y_true = y_true.label_ids
+                    if hasattr(y_true, "labels"):
+                        y_true = y_true.labels
+                    if len(y_score[0][1]) == np.max(y_true) + 1:
+                        y_score = y_score[0]
+                    else:
+                        y_score = y_score[1]
+                    y_score = np.argmax(y_score, axis=1)
+                elif y_true is not None and y_score is not None:
+                    pass  # y_true and y_score are provided
+                else:
+                    raise ValueError(
+                        "Please provide the true and predicted values or a dictionary with 'y_true' and 'y_score'."
+                    )
+
+                y_true, y_score = Metric.flatten(y_true, y_score)
+                y_true_mask_idx = np.where(y_true != self.ignore_y)
+                if self.ignore_y is not None:
+                    y_true = y_true[y_true_mask_idx]
+                    try:
+                        y_score = y_score[y_true_mask_idx]
+                    except Exception as e:
+                        warnings.warn(str(e))
+                kwargs.update(self.kwargs)
+
+                return {name: self.compute(y_true, y_score, *args, **kwargs)}
+
+            return wrapper
+        else:
+            return super().__getattribute__(name)
+
+    def compute(self, y_true, y_score, *args, **kwargs):
+        """
+        Compute the metric, based on the true and predicted values.
+        
+        Args:
+            y_true: The true values
+            y_score: The predicted values
+            *args: Additional positional arguments for the metric
+            **kwargs: Additional keyword arguments for the metric
+            
+        Returns:
+            The computed metric value
+            
+        Raises:
+            NotImplementedError: If no metric function is provided and compute is not implemented
+        """
+        if self.metric_func is not None:
+            kwargs.update(self.kwargs)
+            return self.metric_func(y_true, y_score, *args, **kwargs)
+
+        else:
+            raise NotImplementedError(
+                "Method compute() is not implemented in the child class."
+            )