fusion-bench 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl

fusion_bench/taskpool/dummy.py

@@ -1,5 +1,10 @@
 """
-This is the dummy task pool that is used for debugging purposes.
+Dummy task pool implementation for debugging and testing purposes.
+
+This module provides a minimal task pool implementation that can be used for
+debugging model fusion workflows, testing infrastructure, and validating model
+architectures without running expensive evaluation procedures. It's particularly
+useful during development and prototyping phases.
 """
 
 from typing import Optional
@@ -14,14 +19,41 @@ from fusion_bench.utils.parameters import count_parameters, print_parameters
 
 
 def get_model_summary(model: nn.Module) -> dict:
-    """
-    Generate a report for the given model.
+    """Generate a comprehensive summary report for a PyTorch model.
+
+    Analyzes the given model to extract key information about its architecture,
+    parameter count, and training characteristics. This function is useful for
+    model introspection and comparative analysis during model fusion workflows.
+
+    The summary includes both trainable and total parameter counts, which helps
+    in understanding model complexity and memory requirements. The trainable
+    percentage is particularly useful for identifying models with frozen layers
+    or parameter-efficient fine-tuning setups.
 
     Args:
-        model: The model to generate the report for.
+        model: The PyTorch model to analyze. Can be any nn.Module instance
+            including complex models, fusion models, or pre-trained models.
 
     Returns:
-        dict: The generated report.
+        dict: A structured report containing model information:
+            - model_info: Dictionary with parameter statistics
+                - trainable_params: Number of trainable parameters
+                - all_params: Total number of parameters (trainable + frozen)
+                - trainable_percentage: Ratio of trainable to total parameters
+
+    Example:
+        ```python
+        >>> model = MyModel()
+        >>> summary = get_model_summary(model)
+        >>> print(summary)
+        {
+            "model_info": {
+                "trainable_params": 1234567,
+                "all_params": 1234567,
+                "trainable_percentage": 1.0
+            }
+        }
+        ```
     """
     report = {}
     training_params, all_params = count_parameters(model)
@@ -34,21 +66,77 @@ def get_model_summary(model: nn.Module) -> dict:
 
 
 class DummyTaskPool(BaseTaskPool):
+    """A lightweight task pool implementation for debugging and development workflows.
+
+    This dummy task pool provides a minimal evaluation interface that focuses on
+    model introspection rather than task-specific performance evaluation. It's
+    designed for development scenarios where you need to test model fusion
+    pipelines, validate architectures, or debug workflows without the overhead
+    of running actual evaluation tasks.
+
+    The task pool is particularly useful when:
+        - You want to verify model fusion works correctly
+        - You need to check parameter counts after fusion
+        - You're developing new fusion algorithms
+        - You want to test infrastructure without expensive evaluations
+
+    Example:
+        ```python
+        >>> taskpool = DummyTaskPool(model_save_path="/tmp/fused_model")
+        >>> results = taskpool.evaluate(fused_model)
+        >>> print(f"Model has {results['model_info']['trainable_params']} parameters")
+        ```
     """
-    This is a dummy task pool used for debugging purposes. It inherits from the base TaskPool class.
-    """
 
-    def __init__(self, model_save_path: Optional[str] = None):
-        super().__init__()
+    def __init__(self, model_save_path: Optional[str] = None, **kwargs):
+        """Initialize the dummy task pool with optional model saving capability.
+
+        Args:
+            model_save_path: Optional path where the evaluated model should be saved.
+                If provided, the model will be serialized and saved to this location
+                after evaluation using the separate_save utility. If None, no model
+                saving will be performed.
+
+        Example:
+            ```python
+            >>> # Create taskpool without saving
+            >>> taskpool = DummyTaskPool()
+
+            >>> # Create taskpool with model saving
+            >>> taskpool = DummyTaskPool(model_save_path="/path/to/save/model.pth")
+            ```
+        """
+        super().__init__(**kwargs)
         self.model_save_path = model_save_path
 
     def evaluate(self, model):
-        """
-        Evaluate the given model.
-        This method does nothing but print the parameters of the model in a human-readable format.
+        """Perform lightweight evaluation and analysis of the given model.
+
+        This method provides a minimal evaluation that focuses on model introspection
+        rather than task-specific performance metrics. It performs parameter analysis,
+        optionally saves the model, and returns a summary report.
+
+        The evaluation process includes:
+        1. Printing human-readable parameter information (rank-zero only)
+        2. Optionally saving the model if a save path was configured
+        3. Generating and returning a model summary report
 
         Args:
-            model: The model to evaluate.
+            model: The model to evaluate. Can be any PyTorch nn.Module including
+                fusion models, pre-trained models, or custom architectures.
+
+        Returns:
+            dict: A model summary report containing parameter statistics and
+                architecture information. See get_model_summary() for detailed
+                format specification.
+
+        Example:
+            ```python
+            >>> taskpool = DummyTaskPool(model_save_path="/tmp/model.pth")
+            >>> model = torch.nn.Linear(10, 5)
+            >>> results = taskpool.evaluate(model)
+            >>> print(f"Trainable params: {results['model_info']['trainable_params']}")
+            ```
         """
         if rank_zero_only.rank == 0:
             print_parameters(model, is_human_readable=True)

fusion_bench/taskpool/lm_eval_harness/taskpool.py

@@ -16,6 +16,47 @@ log = logging.getLogger(__name__)
 
 
 class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
+    """A task pool implementation that interfaces with the LM Evaluation Harness framework.
+
+    This class provides a wrapper around the LM Evaluation Harness (lm-eval) library,
+    enabling evaluation of language models on various standardized benchmarks and tasks.
+    It inherits from BaseTaskPool and LightningFabricMixin to provide distributed
+    computing capabilities through PyTorch Lightning Fabric.
+
+    The task pool supports evaluation on multiple tasks simultaneously and provides
+    flexible configuration options for batch processing, output formatting, and
+    logging. It automatically handles model setup and wrapping for distributed
+    evaluation when using Lightning Fabric.
+
+    Args:
+        tasks: A single task name or list of task names to evaluate on.
+            Examples: "hellaswag", ["arc_easy", "arc_challenge", "hellaswag"]
+        apply_chat_template: Whether to apply chat template formatting to inputs.
+            Useful for instruction-tuned or chat models.
+        include_path: Path to additional task definitions or custom tasks.
+        batch_size: Number of samples to process in each batch. Larger values
+            may improve throughput but require more memory.
+        metadata: Additional metadata to include in evaluation results.
+        verbosity: Logging verbosity level for the evaluation process.
+        output_path: Custom path for saving evaluation results. If None,
+            results are saved to the default log directory.
+        log_samples: Whether to log individual sample predictions and targets.
+            Useful for debugging but increases output size significantly.
+        _usage_: Internal usage tracking string.
+        _version_: Internal version tracking string.
+        **kwargs: Additional arguments passed to the LM Evaluation Harness.
+
+    Example:
+        ```python
+        >>> taskpool = LMEvalHarnessTaskPool(
+        ...     tasks=["arc_easy", "hellaswag"],
+        ...     batch_size=8,
+        ...     verbosity="INFO"
+        ... )
+        >>> results = taskpool.evaluate(model)
+        ```
+    """
+
     def __init__(
         self,
         tasks: Union[str, List[str]],
@@ -44,6 +85,45 @@ class LMEvalHarnessTaskPool(BaseTaskPool, LightningFabricMixin):
         self.log_samples = log_samples
 
     def evaluate(self, model, *command_line_args, **kwargs):
+        """Evaluate a language model on the configured tasks using LM Evaluation Harness.
+
+        This method wraps the model with the LM Evaluation Harness framework and
+        executes evaluation on all configured tasks. It automatically handles
+        command-line argument construction, model wrapping with Lightning Fabric
+        for distributed evaluation, and result logging.
+
+        The evaluation process includes:
+        1. Building command-line arguments from instance configuration
+        2. Setting up the LM Evaluation Harness argument parser
+        3. Wrapping the model with Lightning Fabric if not already wrapped
+        4. Creating an HFLM (Hugging Face Language Model) wrapper
+        5. Executing the evaluation through the LM-Eval CLI interface
+
+        Args:
+            model: The language model to evaluate. Can be a Hugging Face model,
+                PyTorch model, or any model compatible with the LM Evaluation Harness.
+                The model will be automatically wrapped with Lightning Fabric for
+                distributed evaluation if not already wrapped.
+            *command_line_args: Additional positional command-line arguments
+                (currently unused but preserved for interface compatibility).
+            **kwargs: Additional keyword arguments that will be converted to
+                command-line flags and passed to the LM Evaluation Harness.
+                Keys will be prefixed with '--' and values converted to strings.
+
+        Returns:
+            None: Results are written to the configured output path and logged.
+
+        Example:
+            ```python
+            >>> taskpool = LMEvalHarnessTaskPool(tasks=["arc_easy"])
+            >>> taskpool.evaluate(model, limit=100, device="cuda")
+            ```
+
+        Note:
+            The method leverages the LM Evaluation Harness's command-line interface
+            internally, which provides standardized evaluation procedures and
+            ensures compatibility with the broader evaluation ecosystem.
+        """
         command_line_args = []
         if self.include_path is not None:
             command_line_args.extend(["--include_path", self.include_path])

fusion_bench/taskpool/nyuv2_taskpool.py

@@ -15,9 +15,37 @@ log = logging.getLogger(__name__)
 
 
 class NYUv2TaskPool(TaskPool):
+    """Task pool for multi-task learning evaluation on the NYUv2 dataset.
+
+    This task pool provides evaluation capabilities for multi-task learning models
+    on the NYU Depth V2 (NYUv2) dataset, which is a popular benchmark for indoor
+    scene understanding. The dataset supports multiple computer vision tasks
+    including semantic segmentation, depth estimation, and surface normal prediction.
+
+    The task pool is designed to work with encoder-decoder architectures where
+    a shared encoder processes input images and task-specific decoders generate
+    predictions for different tasks. It integrates with PyTorch Lightning for
+    streamlined training and evaluation workflows.
+
+    Supported Tasks:
+        - Semantic segmentation
+        - Depth estimation
+        - Surface normal prediction
+    """
+
     _trainer: L.Trainer = None
 
     def __init__(self, taskpool_config: DictConfig):
+        """Initialize the NYUv2 task pool with configuration settings.
+
+        Args:
+            taskpool_config: Configuration object containing all necessary
+                parameters for the task pool, including:
+                - data_dir: Path to the directory containing NYUv2 dataset
+                - tasks: List of tasks to evaluate (e.g., ["semantic", "depth"])
+                - batch_size: Batch size for evaluation data loader
+                - num_workers: Number of worker processes for data loading
+        """
         self.config = taskpool_config
 
     def load_datasets(self):

fusion_bench/utils/__init__.py

@@ -20,3 +20,4 @@ from .packages import import_object
 from .parameters import *
 from .pylogger import get_rankzero_logger
 from .timer import timeit_context
+from .type import BoolStateDictType, StateDictType, TorchModelType

fusion_bench/utils/data.py

@@ -1,6 +1,6 @@
 import pickle
 from pathlib import Path
-from typing import Literal, Optional, Union
+from typing import Any, Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -37,7 +37,9 @@ class InfiniteDataLoader:
         return data
 
 
-def load_tensor_from_file(file_path: Union[str, Path], device=None) -> torch.Tensor:
+def load_tensor_from_file(
+    file_path: Union[str, Path], device: Optional[Union[str, torch.device]] = None
+) -> torch.Tensor:
     """
     Loads a tensor from a file, which can be either a .pt, .pth or .np file.
     If the file is not one of these formats, it will try to load it as a pickle file.
@@ -72,7 +74,7 @@ def train_validation_split(
     validation_size: Optional[int] = None,
     random_seed: Optional[int] = None,
     return_split: Literal["all", "train", "val"] = "both",
-):
+) -> Union[Tuple[Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training and validation set.
 
@@ -134,7 +136,7 @@ def train_validation_test_split(
     test_fraction: float,
     random_seed: Optional[int] = None,
     return_spilt: Literal["all", "train", "val", "test"] = "all",
-):
+) -> Union[Tuple[Dataset, Dataset, Dataset], Dataset]:
     """
     Split a dataset into a training, validation and test set.
 

fusion_bench/utils/devices.py

@@ -1,7 +1,7 @@
 import gc
 import logging
 import os
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 
 import torch
 from transformers.utils import (
@@ -12,6 +12,8 @@ from transformers.utils import (
     is_torch_xpu_available,
 )
 
+from .type import T
+
 __all__ = [
     "clear_cuda_cache",
     "to_device",
@@ -37,7 +39,12 @@ def clear_cuda_cache():
         log.warning("CUDA is not available. No cache to clear.")
 
 
-def to_device(obj, device: Optional[torch.device], **kwargs):
+def to_device(
+    obj: T,
+    device: Optional[torch.device],
+    copy_on_move: bool = False,
+    **kwargs: Any,
+) -> T:
     """
     Move a given object to the specified device.
 
@@ -47,12 +54,20 @@ def to_device(obj, device: Optional[torch.device], **kwargs):
     Args:
         obj: The object to be moved to the device. This can be a torch.Tensor, torch.nn.Module, list, tuple, or dict.
         device (torch.device): The target device to move the object to. This can be `None`.
-        **kwargs: Additional keyword arguments to be passed to the `to` method of torch.Tensor or torch.nn.Module. For example, `non_blocking=True`, `dtype=torch.float16`.
+        copy_on_move (bool, optional): Whether to force a copy operation when moving tensors to a different device.
+            If True, tensors will be copied when moved to a different device (copy=True is passed to tensor.to()).
+            If False (default), tensors are moved without forcing a copy operation, allowing PyTorch to optimize
+            the operation. This parameter only affects torch.Tensor objects; modules and other types are unaffected.
+            Defaults to False.
+        **kwargs: Additional keyword arguments to be passed to the `to` method of torch.Tensor or torch.nn.Module.
+            For example, `non_blocking=True`, `dtype=torch.float16`. Note that if `copy_on_move=True`, the `copy`
+            keyword argument will be automatically set and should not be provided manually.
 
     Returns:
         The object moved to the specified device. The type of the returned object matches the type of the input object.
 
     Examples:
+        ```python
         >>> tensor = torch.tensor([1, 2, 3])
         >>> to_device(tensor, torch.device('cuda'))
         tensor([1, 2, 3], device='cuda:0')
@@ -64,17 +79,26 @@ def to_device(obj, device: Optional[torch.device], **kwargs):
         >>> data = [torch.tensor([1, 2]), torch.tensor([3, 4])]
         >>> to_device(data, torch.device('cuda'))
         [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')]
+
+        >>> # Force copy when moving to different device
+        >>> tensor = torch.tensor([1, 2, 3], device='cpu')
+        >>> copied_tensor = to_device(tensor, torch.device('cuda'), copy_on_move=True)
+        >>> # tensor and copied_tensor will have different memory locations
+        ```
     """
-    if isinstance(obj, (torch.Tensor, torch.nn.Module)):
+    if isinstance(obj, torch.Tensor):
+        if copy_on_move:
+            if obj.device != torch.device(device):
+                kwargs["copy"] = True
+        return obj.to(device, **kwargs)
+    elif isinstance(obj, torch.nn.Module):
         return obj.to(device, **kwargs)
     elif isinstance(obj, list):
-        return [to_device(o, device) for o in obj]
+        return [to_device(o, device, **kwargs) for o in obj]
     elif isinstance(obj, tuple):
-        return tuple(to_device(o, device) for o in obj)
+        return tuple(to_device(o, device, **kwargs) for o in obj)
     elif isinstance(obj, dict):
-        for key in obj:
-            obj[key] = to_device(obj[key], device)
-        return obj
+        return {key: to_device(value, device, **kwargs) for key, value in obj.items()}
     else:
         # the default behavior is to return the object as is
         return obj
@@ -102,7 +126,7 @@ def num_devices(devices: Union[int, List[int], str]) -> int:
         )
 
 
-def get_device(obj) -> torch.device:
+def get_device(obj: Any) -> torch.device:
     """
     Get the device of a given object.
 
@@ -151,6 +175,7 @@ def get_current_device() -> torch.device:
                     If not set, it defaults to "0".
 
     Example:
+
         >>> device = get_current_device()
         >>> print(device)
         xpu:0  # or npu:0, mps:0, cuda:0, cpu depending on availability
@@ -241,7 +266,7 @@ def cleanup_cuda():
     torch.cuda.reset_peak_memory_stats()
 
 
-def print_memory_usage(print_fn=print):
+def print_memory_usage(print_fn=print) -> str:
     """
     Print the current GPU memory usage.
 

fusion_bench/utils/dtype.py

@@ -1,5 +1,5 @@
 import contextlib
-from typing import Dict, Generator, Iterable, Optional, Tuple
+from typing import Dict, Generator, Iterable, Optional, Tuple, Union
 
 import torch
 from transformers.utils import (
@@ -25,7 +25,7 @@ PRECISION_STR_TO_DTYPE: Dict[str, torch.dtype] = {
 }
 
 
-def parse_dtype(dtype: Optional[str]):
+def parse_dtype(dtype: Optional[str]) -> Optional[torch.dtype]:
     """
     Parses a string representation of a data type and returns the corresponding torch.dtype.
 
@@ -92,6 +92,7 @@ def set_default_dtype(dtype: torch.dtype) -> Generator[None, None, None]:
         ContextManager: context manager for setting default dtype.
 
     Example:
+
         >>> with set_default_dtype(torch.bfloat16):
         >>>     x = torch.tensor([1, 2, 3])
         >>>     x.dtype

fusion_bench/utils/lazy_state_dict.py

@@ -2,7 +2,18 @@ import json
 import logging
 import os
 from copy import deepcopy
-from typing import TYPE_CHECKING, Dict, Iterator, List, Mapping, Optional, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Generic,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 import torch
 from accelerate import init_empty_weights
@@ -11,10 +22,12 @@ from huggingface_hub import snapshot_download
 from safetensors import safe_open
 from safetensors.torch import load_file
 from torch import nn
+from torch.nn.modules.module import _IncompatibleKeys
 from transformers import AutoConfig
 
 from fusion_bench.utils.dtype import parse_dtype
 from fusion_bench.utils.packages import import_object
+from fusion_bench.utils.type import TorchModelType
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
@@ -49,7 +62,7 @@ def resolve_checkpoint_path(
         )
 
 
-class LazyStateDict(Mapping[str, torch.Tensor]):
+class LazyStateDict(Mapping[str, torch.Tensor], Generic[TorchModelType]):
     """
     Dictionary-like object that lazily loads a state dict from a checkpoint path.
     """
@@ -63,11 +76,14 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
     _index: Optional[Dict[str, str]]
     """Mapping of parameter names to checkpoint files."""
 
+    meta_module: TorchModelType = None
+    meta_module_class: Optional[Type[TorchModelType]] = None
+
     def __init__(
         self,
         checkpoint: str,
-        meta_module_class: Optional[Type[nn.Module]] = None,
-        meta_module: Optional[nn.Module] = None,
+        meta_module_class: Optional[Type[TorchModelType]] = None,
+        meta_module: Optional[TorchModelType] = None,
         cache_state_dict: bool = False,
         torch_dtype: Optional[torch.dtype] = None,
         device: str = "cpu",
@@ -88,15 +104,19 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
             hf_proxies (Dict, optional): Proxies to use for downloading from Hugging Face Hub.
         """
         self.cache_state_dict = cache_state_dict
+
+        # Validate that both meta_module_class and meta_module are not provided
+        if meta_module_class is not None and meta_module is not None:
+            raise ValueError(
+                "Cannot provide both meta_module_class and meta_module, please provide only one."
+            )
+
         self.meta_module_class = meta_module_class
         if isinstance(self.meta_module_class, str):
             self.meta_module_class = import_object(self.meta_module_class)
         self.meta_module = meta_module
+
         if self.meta_module_class is not None:
-            if self.meta_module is not None:
-                raise ValueError(
-                    "Cannot provide both meta_module_class and meta_module, please provide only one."
-                )
             with init_empty_weights():
                 self.meta_module = self.meta_module_class.from_pretrained(
                     checkpoint,
@@ -173,9 +193,13 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
         """
         `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
         """
+        if hasattr(self, "_cached_dtype"):
+            return self._cached_dtype
+
         first_key = next(iter(self.keys()))
         first_param = self[first_key]
-        return first_param.dtype
+        self._cached_dtype = first_param.dtype
+        return self._cached_dtype
 
     def state_dict(self, keep_vars: bool = False) -> "LazyStateDict":
         """
@@ -321,9 +345,7 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
         if self._state_dict_cache is not None:
             self._state_dict_cache[key] = value
         else:
-            log.warning(
-                "State dict cache is disabled, setting a tensor will not update the cache."
-            )
+            log.warning("State dict cache is disabled, initializing the cache.")
             self._state_dict_cache = {key: value}
 
     def __contains__(self, key: str) -> bool:
@@ -339,7 +361,7 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
                     self._checkpoint_files[0], key, update_cache=False
                 )
                 return tensor is not None
-            except Exception:
+            except (KeyError, FileNotFoundError, RuntimeError, EOFError):
                 return False
         return False
 
@@ -409,8 +431,8 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
             )
 
     def load_state_dict(
-        self, state_dict: Dict[str, torch.Tensor], strict: bool = True
-    ) -> None:
+        self, state_dict: Mapping[str, torch.Tensor], strict: bool = True
+    ) -> _IncompatibleKeys:
         """
         Load a state dict into this LazyStateDict.
         This method is only for compatibility with nn.Module and it overrides the cache of LazyStateDict.
@@ -419,16 +441,60 @@ class LazyStateDict(Mapping[str, torch.Tensor]):
             state_dict (Dict[str, torch.Tensor]): The state dict to load.
             strict (bool): Whether to enforce that all keys in the state dict are present in this LazyStateDict.
         """
+        if not isinstance(state_dict, Mapping):
+            raise TypeError(
+                f"Expected state_dict to be dict-like, got {type(state_dict)}."
+            )
+
+        missing_keys: list[str] = []
+        unexpected_keys: list[str] = []
+        error_msgs: list[str] = []
+
         log.warning(
             "Loading state dict into LazyStateDict is not recommended, as it may lead to unexpected behavior. "
             "Use with caution."
         )
+
+        # Check for unexpected keys in the provided state_dict
+        for key in state_dict:
+            if key not in self:
+                unexpected_keys.append(key)
+
+        # Check for missing keys that are expected in this LazyStateDict
+        for key in self.keys():
+            if key not in state_dict:
+                missing_keys.append(key)
+
+        # Handle strict mode
         if strict:
-            for key in state_dict:
-                if key not in self:
-                    raise KeyError(f"Key {key} not found in LazyStateDict.")
+            if len(unexpected_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Unexpected key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in unexpected_keys)
+                    ),
+                )
+            if len(missing_keys) > 0:
+                error_msgs.insert(
+                    0,
+                    "Missing key(s) in state_dict: {}. ".format(
+                        ", ".join(f'"{k}"' for k in missing_keys)
+                    ),
+                )
+
+        if len(error_msgs) > 0:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    self.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+
+        # Load the state dict values
         for key, value in state_dict.items():
-            self[key] = value
+            if key in self:  # Only set keys that exist in this LazyStateDict
+                self[key] = value
+
+        return _IncompatibleKeys(missing_keys, unexpected_keys)
 
     def __getattr__(self, name: str):
         if "meta_module" in self.__dict__:

fusion_bench/utils/packages.py

@@ -1,7 +1,7 @@
 import importlib.metadata
 import importlib.util
 from functools import lru_cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from packaging import version
 
@@ -69,7 +69,7 @@ def is_vllm_available():
     return _is_package_available("vllm")
 
 
-def import_object(abs_obj_name: str):
+def import_object(abs_obj_name: str) -> Any:
     """
     Imports a class from a module given the absolute class name.
 
@@ -84,7 +84,7 @@ def import_object(abs_obj_name: str):
     return getattr(module, obj_name)
 
 
-def compare_versions(v1, v2):
+def compare_versions(v1: str, v2: str) -> int:
     """Compare two version strings.
     Returns -1 if v1 < v2, 0 if v1 == v2, 1 if v1 > v2"""