Modernize type hints - dataclass arguments and recipe metadata

Signed-off-by: ojeda-e <[email protected]>

Author: ojeda-e

src/llmcompressor/args/dataset_arguments.py

@@ -8,7 +8,7 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable
 
 from transformers import DefaultDataCollator
 
@@ -19,7 +19,7 @@ class DVCDatasetArguments:
     Arguments for training using DVC
     """
 
-    dvc_data_repository: Optional[str] = field(
+    dvc_data_repository: str | None = field(
         default=None,
         metadata={"help": "Path to repository used for dvc_dataset_path"},
     )
@@ -31,7 +31,7 @@ class CustomDatasetArguments(DVCDatasetArguments):
     Arguments for training using custom datasets
     """
 
-    dataset_path: Optional[str] = field(
+    dataset_path: str | None = field(
         default=None,
         metadata={
             "help": (
@@ -52,12 +52,12 @@ class CustomDatasetArguments(DVCDatasetArguments):
         },
     )
 
-    remove_columns: Union[None, str, List] = field(
+    remove_columns: str | list[str] | None = field(
         default=None,
         metadata={"help": "Column names to remove after preprocessing (deprecated)"},
     )
 
-    preprocessing_func: Union[None, str, Callable] = field(
+    preprocessing_func: str | Callable | None = field(
         default=None,
         metadata={
             "help": (
@@ -85,7 +85,7 @@ class DatasetArguments(CustomDatasetArguments):
     arguments to be able to specify them on the command line
     """
 
-    dataset: Optional[str] = field(
+    dataset: str | None = field(
         default=None,
         metadata={
             "help": (
@@ -94,7 +94,7 @@ class DatasetArguments(CustomDatasetArguments):
             )
         },
     )
-    dataset_config_name: Optional[str] = field(
+    dataset_config_name: str | None = field(
         default=None,
         metadata={
             "help": ("The configuration name of the dataset to use"),
@@ -114,15 +114,15 @@ class DatasetArguments(CustomDatasetArguments):
             "help": "Whether or not to concatenate datapoints to fill max_seq_length"
         },
     )
-    raw_kwargs: Dict = field(
+    raw_kwargs: dict = field(
         default_factory=dict,
         metadata={"help": "Additional keyboard args to pass to datasets load_data"},
     )
-    splits: Union[None, str, List, Dict] = field(
+    splits: str | list[str] | dict[str, str] | None = field(
         default=None,
         metadata={"help": "Optional percentages of each split to download"},
     )
-    num_calibration_samples: Optional[int] = field(
+    num_calibration_samples: int | None = field(
         default=512,
         metadata={"help": "Number of samples to use for one-shot calibration"},
     )
@@ -136,21 +136,21 @@ class DatasetArguments(CustomDatasetArguments):
             "module definitions"
         },
     )
-    shuffle_calibration_samples: Optional[bool] = field(
+    shuffle_calibration_samples: bool | None = field(
         default=True,
         metadata={
             "help": "whether to shuffle the dataset before selecting calibration data"
         },
     )
-    streaming: Optional[bool] = field(
+    streaming: bool | None = field(
         default=False,
         metadata={"help": "True to stream data from a cloud dataset"},
     )
     overwrite_cache: bool = field(
         default=False,
         metadata={"help": "Overwrite the cached preprocessed datasets or not."},
     )
-    preprocessing_num_workers: Optional[int] = field(
+    preprocessing_num_workers: int | None = field(
         default=None,
         metadata={"help": "The number of processes to use for the preprocessing."},
     )
@@ -162,14 +162,14 @@ class DatasetArguments(CustomDatasetArguments):
             "in the batch (which can be faster on GPU but will be slower on TPU)."
         },
     )
-    max_train_samples: Optional[int] = field(
+    max_train_samples: int | None = field(
         default=None,
         metadata={
             "help": "For debugging purposes or quicker training, truncate the number "
             "of training examples to this value if set."
         },
     )
-    min_tokens_per_module: Optional[float] = field(
+    min_tokens_per_module: float | None = field(
         default=None,
         metadata={
             "help": (
@@ -182,15 +182,15 @@ class DatasetArguments(CustomDatasetArguments):
         },
     )
     # --- pipeline arguments --- #
-    pipeline: Optional[str] = field(
+    pipeline: str | None = field(
         default="independent",
         metadata={
             "help": "Calibration pipeline used to calibrate model"
             "Options: ['basic', 'datafree', 'sequential', 'layer_sequential', "
             "independent]"
         },
     )
-    tracing_ignore: List[str] = field(
+    tracing_ignore: list[str] = field(
         default_factory=lambda: [
             "_update_causal_mask",
             "create_causal_mask",
@@ -209,7 +209,7 @@ class DatasetArguments(CustomDatasetArguments):
             "{module}.{method_name} or {function_name}"
         },
     )
-    sequential_targets: Optional[List[str]] = field(
+    sequential_targets: list[str] | None = field(
         default=None,
         metadata={
             "help": "List of layer targets for the sequential pipeline. "

src/llmcompressor/args/model_arguments.py

@@ -8,7 +8,6 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Optional
 
 
 @dataclass
@@ -27,31 +26,31 @@ class ModelArguments:
             )
         },
     )
-    distill_teacher: Optional[str] = field(
+    distill_teacher: str | None = field(
         default=None,
         metadata={
             "help": "Teacher model (a trained text generation model)",
         },
     )
-    config_name: Optional[str] = field(
+    config_name: str | None = field(
         default=None,
         metadata={
             "help": "Pretrained config name or path if not the same as model_name"
         },
     )
-    tokenizer: Optional[str] = field(
+    tokenizer: str | None = field(
         default=None,
         metadata={
             "help": "Pretrained tokenizer name or path if not the same as model_name"
         },
     )
-    processor: Optional[str] = field(
+    processor: str | None = field(
         default=None,
         metadata={
             "help": "Pretrained processor name or path if not the same as model_name"
         },
     )
-    cache_dir: Optional[str] = field(
+    cache_dir: str | None = field(
         default=None,
         metadata={"help": "Where to store the pretrained data from huggingface.co"},
     )
@@ -85,7 +84,7 @@ class ModelArguments:
         },
     )
     # TODO: potentialy separate out/expand to additional saving args
-    save_compressed: Optional[bool] = field(
+    save_compressed: bool | None = field(
         default=True,
         metadata={"help": "Whether to compress sparse models during save"},
     )

src/llmcompressor/args/recipe_arguments.py

@@ -7,20 +7,19 @@
 """
 
 from dataclasses import dataclass, field
-from typing import List, Optional
 
 
 @dataclass
 class RecipeArguments:
     """Recipe and session variables"""
 
-    recipe: Optional[str] = field(
+    recipe: str | None = field(
         default=None,
         metadata={
             "help": "Path to a LLM Compressor sparsification recipe",
         },
     )
-    recipe_args: Optional[List[str]] = field(
+    recipe_args: list[str] | None = field(
         default=None,
         metadata={
             "help": (
@@ -29,7 +28,7 @@ class RecipeArguments:
             )
         },
     )
-    clear_sparse_session: Optional[bool] = field(
+    clear_sparse_session: bool | None = field(
         default=False,
         metadata={
             "help": (
@@ -38,7 +37,7 @@ class RecipeArguments:
             )
         },
     )
-    stage: Optional[str] = field(
+    stage: str | None = field(
         default=None,
         metadata={"help": ("The stage of the recipe to use for oneshot / train.",)},
     )

src/llmcompressor/args/training_arguments.py

@@ -8,7 +8,6 @@
 """
 
 from dataclasses import dataclass, field
-from typing import Optional
 
 from transformers import TrainingArguments as HFTrainingArgs
 
@@ -25,11 +24,11 @@ class TrainingArguments(HFTrainingArgs):
 
     """
 
-    do_oneshot: Optional[bool] = field(
+    do_oneshot: bool | None = field(
         default=False,
         metadata={"help": "Whether to run one-shot calibration in stages"},
     )
-    run_stages: Optional[bool] = field(
+    run_stages: bool | None = field(
         default=False, metadata={"help": "Whether to trigger recipe stage by stage"}
     )
     output_dir: str = field(

src/llmcompressor/args/utils.py

@@ -7,8 +7,6 @@
 warnings, and processor resolution.
 """
 
-from typing import Tuple
-
 from loguru import logger
 from transformers import HfArgumentParser
 
@@ -23,7 +21,13 @@
 
 def parse_args(
     include_training_args: bool = False, **kwargs
-) -> Tuple[ModelArguments, DatasetArguments, RecipeArguments, TrainingArguments, str]:
+) -> tuple[
+    ModelArguments,
+    DatasetArguments,
+    RecipeArguments,
+    TrainingArguments | None,
+    str | None,
+]:
     """
     Keyword arguments passed in from `oneshot` or `train` will
     separate the arguments into the following:

src/llmcompressor/recipe/metadata.py

@@ -6,7 +6,7 @@
 structured data containers for recipe configuration and execution tracking.
 """
 
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from pydantic import BaseModel, Field
 
@@ -22,7 +22,7 @@ class DatasetMetaData(BaseModel):
     name: str = None
     version: str = None
     hash: str = None
-    shape: List[int] = Field(default_factory=list)
+    shape: list[int] = Field(default_factory=list)
     num_classes: int = None
     num_train_samples: int = None
     num_val_samples: int = None
@@ -31,24 +31,24 @@ class DatasetMetaData(BaseModel):
 
 class ParamMetaData(BaseModel):
     name: str = None
-    shape: List[int] = None
+    shape: list[int] = None
     weight_hash: str = None
 
 
 class LayerMetaData(BaseModel):
     name: str = None
     type: str = None
     index: int = None
-    attributes: Dict[str, Any] = None
-    input_shapes: List[List[int]] = None
-    output_shapes: List[List[int]] = None
-    params: Dict[str, ParamMetaData] = None
+    attributes: dict[str, Any] = None
+    input_shapes: list[list[int]] = None
+    output_shapes: list[list[int]] = None
+    params: dict[str, ParamMetaData] = None
 
 
 class ModelMetaData(BaseModel):
     architecture: str = None
     sub_architecture: str = None
-    input_shapes: List[List[int]] = None
-    output_shapes: List[List[int]] = None
-    layers: List[LayerMetaData] = Field(default_factory=list)
-    layer_prefix: Optional[str] = None
+    input_shapes: list[list[int]] = None
+    output_shapes: list[list[int]] = None
+    layers: list[LayerMetaData] = Field(default_factory=list)
+    layer_prefix: str | None = None