Modernize transformers module with type hints and generic types

Author: sugatmahanti

src/llmcompressor/transformers/compression/compressed_tensors_utils.py

@@ -2,7 +2,6 @@
 import weakref
 from collections.abc import Generator
 from functools import wraps
-from typing import Optional
 
 import torch
 from accelerate.accelerator import get_state_dict_offloaded_model
@@ -54,8 +53,8 @@ def save_pretrained_compressed(save_pretrained_method):
         @wraps(original_save_pretrained)
         def save_pretrained_wrapper(
             save_directory: str,
-            sparsity_config: Optional[SparsityCompressionConfig] = None,
-            quantization_format: Optional[str] = None,
+            sparsity_config: SparsityCompressionConfig | None = None,
+            quantization_format: str | None = None,
             save_compressed: bool = True,
             safe_serialization: bool = True,
             skip_sparsity_compression_stats: bool = True,
@@ -233,8 +232,8 @@ def untie_if_target_shared_embedding(
 
 def get_model_compressor(
     model: torch.nn.Module,
-    sparsity_config: Optional[SparsityCompressionConfig] = None,
-    quantization_format: Optional[str] = None,
+    sparsity_config: SparsityCompressionConfig | None = None,
+    quantization_format: str | None = None,
     save_compressed: bool = True,
     skip_sparsity_compression_stats: bool = True,
     disable_sparse_compression: bool = False,

src/llmcompressor/transformers/compression/helpers.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Dict, List, Optional, Tuple
+from typing import Tuple
 
 import torch
 from accelerate.accelerator import get_state_dict_offloaded_model
@@ -51,8 +51,8 @@ def tensor_follows_mask_structure(tensor: torch.Tensor, mask: str = "2:4") -> bo
 
 
 def infer_sparsity_structure_from_modifiers(
-    modifiers: List[Modifier],  # noqa E501
-) -> Optional[str]:
+    modifiers: list[Modifier],  # noqa E501
+) -> str | None:
     """
     Determines the sparsity structure, if any exists, given the list of modifiers.
 
@@ -65,7 +65,7 @@ def infer_sparsity_structure_from_modifiers(
     return None
 
 
-def infer_sparsity_structure_from_model(model: torch.nn.Module) -> Optional[str]:
+def infer_sparsity_structure_from_model(model: torch.nn.Module) -> str | None:
     """
     Determines the sparsity structure, if any exists, given the model
 
@@ -104,7 +104,7 @@ def infer_sparse_targets_and_ignores(
     model: torch.nn.Module,
     sparsity_structure: str,
     sparsity_threshold: float,
-) -> Tuple[List[str], List[str]]:
+) -> Tuple[list[str], list[str]]:
     """
     Infers the target and ignore layers in the given model
     to be used for sparsity compression
@@ -151,7 +151,7 @@ def is_sparse_compression_target(
 
 def _get_sparse_targets_ignore_dicts(
     module: torch.nn.Module, sparsity_structure: str, sparsity_threshold: float
-) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]:
+) -> Tuple[dict[str, list[str]], dict[str, list[str]]]:
     """
     Get sparse targets and ignore dictionaries
 
@@ -176,8 +176,8 @@ def _get_sparse_targets_ignore_dicts(
 
 
 def _reduce_targets_and_ignores_into_lists(
-    exhaustive_targets: Dict[str, List[str]], exhaustive_ignore: Dict[str, List[str]]
-) -> Tuple[List[str], List[str]]:
+    exhaustive_targets: dict[str, list[str]], exhaustive_ignore: dict[str, list[str]]
+) -> Tuple[list[str], list[str]]:
     """
     Reduces the targets and ignores dictionaries into lists
 

src/llmcompressor/transformers/compression/sparsity_metadata_config.py

@@ -1,5 +1,3 @@
-from typing import Dict, List, Optional
-
 from compressed_tensors import CompressionFormat, SparsityCompressionConfig
 from compressed_tensors.config import SparsityStructure
 from compressed_tensors.quantization import QuantizationType
@@ -30,7 +28,7 @@ class SparsityConfigMetadata:
 
     @staticmethod
     def infer_global_sparsity(
-        model: Module, state_dict: Optional[Dict[str, Tensor]] = None
+        model: Module, state_dict: dict[str, Tensor] | None = None
     ) -> float:
         """
         Calculates the global percentage of sparse zero weights in the model
@@ -47,12 +45,12 @@ def infer_global_sparsity(
 
     @staticmethod
     def infer_sparsity_structure(
-        model: Optional[Module] = None, check_only_modifiers: Optional[bool] = False
+        model: Module | None = None, check_only_modifiers: bool | None = False
     ) -> str:
         """
         Determines what sparsity structure, if any, was applied.
 
-        First, there is an attempt to dedue the sparsity structure
+        First, there is an attempt to deduce the sparsity structure
         from the currently active sparse session.
 
         If that fails, the sparsity structure is inferred from the
@@ -83,12 +81,12 @@ def infer_sparsity_structure(
     @staticmethod
     def from_pretrained(
         model: Module,
-        state_dict: Optional[Dict[str, Tensor]] = None,
+        state_dict: dict[str, Tensor] | None = None,
         compress: bool = False,
-        quantization_format: Optional[CompressionFormat] = None,
+        quantization_format: CompressionFormat | None = None,
         disable_sparse_compression: bool = False,
-        sparsity_structure: Optional[str] = None,
-    ) -> Optional["SparsityCompressionConfig"]:
+        sparsity_structure: str | None = None,
+    ) -> "SparsityCompressionConfig" | None:
         """
         Determines compression type and informational parameters for a given model
 
@@ -155,7 +153,7 @@ def from_pretrained(
     def fill_config_details(
         config: SparsityCompressionConfig,
         model: Module,
-        state_dict: Optional[Dict[str, Tensor]] = None,
+        state_dict: dict[str, Tensor] | None = None,
     ):
         """
         Fills in informational sparsity parameters from a given model
@@ -173,7 +171,7 @@ def fill_config_details(
     @staticmethod
     def is_sparse24_bitmask_supported(
         model: Module,
-        sparsity_structure: Optional[str] = None,
+        sparsity_structure: str | None = None,
     ) -> bool:
         """
         Determines if sparse 24 bitmask sparse compressor is supported for a given model
@@ -202,7 +200,7 @@ def is_sparse24_bitmask_supported(
 
         # when model is quantized, and has 2:4 sparsity
 
-        supported_scheme_types: List[str] = [
+        supported_scheme_types: list[str] = [
             QuantizationType.INT.value,
             QuantizationType.FLOAT.value,
         ]

src/llmcompressor/transformers/finetune/data/base.py

@@ -10,7 +10,7 @@
 import inspect
 from functools import cached_property
 from inspect import _ParameterKind as Kind
-from typing import Any, Callable, Dict, List, Union
+from typing import Any, Callable
 
 from compressed_tensors.registry import RegistryMixin
 from datasets import Dataset, IterableDataset
@@ -202,7 +202,7 @@ def load_dataset(self):
         )
 
     @cached_property
-    def preprocess(self) -> Union[Callable[[LazyRow], Any], None]:
+    def preprocess(self) -> Callable[[LazyRow], Any] | None:
         """
         The function must return keys which correspond to processor/tokenizer kwargs,
         optionally including PROMPT_KEY
@@ -225,7 +225,7 @@ def preprocess(self) -> Union[Callable[[LazyRow], Any], None]:
         return self.dataset_template
 
     @property
-    def dataset_template(self) -> Union[Callable[[Any], Any], None]:
+    def dataset_template(self) -> Callable[[Any], Any] | None:
         return None
 
     def rename_columns(self, dataset: DatasetType) -> DatasetType:
@@ -254,7 +254,7 @@ def filter_tokenizer_args(self, dataset: DatasetType) -> DatasetType:
             list(set(column_names) - set(tokenizer_args) - set([self.PROMPT_KEY]))
         )
 
-    def tokenize(self, data: LazyRow) -> Dict[str, Any]:
+    def tokenize(self, data: LazyRow) -> dict[str, Any]:
         # separate prompt
         prompt = data.pop(self.PROMPT_KEY, None)
 
@@ -276,7 +276,7 @@ def tokenize(self, data: LazyRow) -> Dict[str, Any]:
 
         return data
 
-    def group_text(self, data: LazyRow) -> Dict[str, Any]:
+    def group_text(self, data: LazyRow) -> dict[str, Any]:
         concatenated_data = {k: sum(data[k], []) for k in data.keys()}
         total_length = len(concatenated_data[list(data.keys())[0]])
         total_length = (total_length // self.max_seq_length) * self.max_seq_length
@@ -311,10 +311,10 @@ def add_labels(self, data: LazyRow) -> LazyRow:
 
     def map(
         self,
-        dataset: Union[Dataset, IterableDataset],
+        dataset: Dataset | IterableDataset,
         function: Callable[[Any], Any],
         **kwargs,
-    ) -> Union[Dataset, IterableDataset]:
+    ) -> Dataset | IterableDataset:
         """
         Wrapper function around Dataset.map and IterableDataset.map.
 
@@ -336,7 +336,7 @@ def map(
         return dataset
 
 
-def get_columns(dataset: DatasetType) -> List[str]:
+def get_columns(dataset: DatasetType) -> list[str]:
     column_names = dataset.column_names
     if isinstance(column_names, dict):
         column_names = sum(column_names.values(), [])

src/llmcompressor/transformers/finetune/data/data_helpers.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import Any, Dict, Optional
+from typing import Any
 
 from datasets import Dataset, load_dataset
 
@@ -15,8 +15,8 @@
 
 def get_raw_dataset(
     dataset_args,
-    cache_dir: Optional[str] = None,
-    streaming: Optional[bool] = False,
+    cache_dir: str | None = None,
+    streaming: bool | None = False,
     **kwargs,
 ) -> Dataset:
     """
@@ -37,7 +37,7 @@ def get_raw_dataset(
     return raw_datasets
 
 
-def get_custom_datasets_from_path(path: str, ext: str = "json") -> Dict[str, str]:
+def get_custom_datasets_from_path(path: str, ext: str = "json") -> dict[str, str]:
     """
     Get a dictionary of custom datasets from a directory path. Support HF's load_dataset
      for local folder datasets https://huggingface.co/docs/datasets/loading
@@ -105,7 +105,7 @@ def get_custom_datasets_from_path(path: str, ext: str = "json") -> Dict[str, str
     return transform_dataset_keys(data_files)
 
 
-def transform_dataset_keys(data_files: Dict[str, Any]):
+def transform_dataset_keys(data_files: dict[str, Any]):
     """
     Transform dict keys to `train`, `val` or `test` for the given input dict
     if matches exist with the existing keys. Note that there can only be one

src/llmcompressor/transformers/finetune/data/peoples_speech.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Dict
+from typing import TYPE_CHECKING, Any
 
 from datasets.formatting.formatting import LazyRow
 from loguru import logger
@@ -68,7 +68,7 @@ def filter_tokenizer_args(self, dataset: DatasetType) -> DatasetType:
         else:
             return super().filter_tokenizer_args(dataset)
 
-    def tokenize(self, data: LazyRow) -> Dict[str, Any]:
+    def tokenize(self, data: LazyRow) -> dict[str, Any]:
         if self.processor_type == "WhisperProcessor":
             inputs = self.processor(
                 audio=data["audio"],

src/llmcompressor/transformers/finetune/session_mixin.py

@@ -2,7 +2,7 @@
 import math
 import os
 from dataclasses import asdict
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Tuple
 
 import torch
 from loguru import logger
@@ -56,9 +56,9 @@ def __init__(
         self,
         recipe: str,
         model_args: "ModelArguments",
-        dataset_args: Optional["DatasetArguments"] = None,
-        teacher: Optional[Union[Module, str]] = None,
-        recipe_args: Optional[Union[Dict[str, Any], str]] = None,
+        dataset_args: "DatasetArguments" | None = None,
+        teacher: Module | str | None = None,
+        recipe_args: dict[str, Any] | str | None = None,
         **kwargs,
     ):
         self.recipe = recipe
@@ -125,8 +125,8 @@ def __init__(
     def initialize_session(
         self,
         epoch: float,
-        checkpoint: Optional[str] = None,
-        stage: Optional[str] = None,
+        checkpoint: str | None = None,
+        stage: str | None = None,
     ):
         """
         Initialize the CompressionSession from the specified epoch, evaluates the recipe
@@ -251,8 +251,8 @@ def create_scheduler(
     def training_step(
         self,
         model: torch.nn.Module,
-        inputs: Dict[str, Union[torch.Tensor, Any]],
-        num_items_in_batch: Optional[int] = None,
+        inputs: dict[str, torch.Tensor | Any],
+        num_items_in_batch: int | None = None,
     ) -> torch.Tensor:
         """
         Overrides the Trainer's training step to trigger the batch_start callback to
@@ -274,10 +274,10 @@ def training_step(
     def compute_loss(
         self,
         model: Module,
-        inputs: Dict[str, Any],
+        inputs: dict[str, Any],
         return_outputs: bool = False,
-        num_items_in_batch: Optional[torch.Tensor] = None,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, Any]]:
+        num_items_in_batch: torch.Tensor | None = None,
+    ) ->torch.Tensor | Tuple[torch.Tensor, Any]:
         """
         Override for the compute_loss to factor trigger callbacks and filter columns
 
@@ -326,7 +326,7 @@ def compute_loss(
 
         return loss
 
-    def train(self, *args, stage: Optional[str] = None, **kwargs):
+    def train(self, *args, stage: str | None = None, **kwargs):
         """
         Run a sparsification training cycle. Runs initialization for the sparse session
         before calling super().train() and finalization of the session after.
@@ -370,7 +370,7 @@ def save_model(
         self,
         output_dir: str,
         _internal_call: bool = False,
-        skip_sparsity_compression_stats: Optional[bool] = True,
+        skip_sparsity_compression_stats: bool | None = True,
     ):
         """
         Override of the save_model function and expects it to exist in the parent.
@@ -478,10 +478,10 @@ def _prepare_model_for_fsdp(self):
 
     def _extract_metadata(
         self,
-        metadata_args: List[str],
-        training_args_dict: Dict[str, Any],
-        dataset_args_dict: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        metadata_args: list[str],
+        training_args_dict: dict[str, Any],
+        dataset_args_dict: dict[str, Any],
+    ) -> dict[str, Any]:
         metadata = {}
         if not training_args_dict.keys().isdisjoint(dataset_args_dict.keys()):
             raise ValueError(
@@ -509,7 +509,7 @@ def _check_super_defined(self, func: str):
                 f"The super class for SessionManagerMixIn must define a {func} function"
             )
 
-    def _calculate_checkpoint_info(self, kwargs) -> Tuple[Optional[str], float]:
+    def _calculate_checkpoint_info(self, kwargs) -> Tuple[str | None, float]:
         """
         If resuming from checkpoint is set, get checkpoint and epoch to resume from
         """