style quality

Author: ved1beta

src/llmcompressor/entrypoints/oneshot.py

@@ -1,4 +1,4 @@
-from typing import Optional, List
+from typing import List, Optional
 
 from torch.utils.data import DataLoader
 from transformers import PreTrainedModel
@@ -190,13 +190,11 @@ def oneshot(
     save_compressed: bool = True,
     oneshot_device: str = "cuda:0",
     model_revision: str = "main",
-    
     # Recipe parameters
     recipe: Optional[str] = None,
     recipe_args: Optional[List[str]] = None,
     clear_sparse_session: bool = False,
     stage: Optional[str] = None,
-    
     # Dataset parameters
     dataset: Optional[str] = None,
     dataset_config_name: Optional[str] = None,
@@ -212,65 +210,75 @@ def oneshot(
     preprocessing_num_workers: Optional[int] = None,
     min_tokens_per_module: Optional[float] = None,
     trust_remote_code_data: bool = False,
-    
     # Output parameters
     output_dir: Optional[str] = None,
-    
     # For backward compatibility
-    **kwargs
+    **kwargs,
 ) -> PreTrainedModel:
     """
     Performs oneshot calibration on a model.
-    
+
     Args:
         # Model arguments
-        model (str): A pretrained model identifier from huggingface.co/models or a path to a
-            local model. Required parameter.
-        distill_teacher (Optional[str]): Teacher model (a trained text generation model) for
-            distillation.
-        config_name (Optional[str]): Pretrained config name or path if not the same as model_name.
-        tokenizer (Optional[str]): Pretrained tokenizer name or path if not the same as model_name.
-        processor (Optional[str]): Pretrained processor name or path if not the same as model_name.
-        cache_dir (Optional[str]): Where to store the pretrained data from huggingface.co.
-        use_auth_token (bool): Whether to use Hugging Face auth token for private models.
+        model (str): A pretrained model identifier from huggingface.co/models or a path
+            to a local model. Required parameter.
+        distill_teacher (Optional[str]): Teacher model (a trained text generation model)
+            for distillation.
+        config_name (Optional[str]): Pretrained config name or path if not the same as
+            model_name.
+        tokenizer (Optional[str]): Pretrained tokenizer name or path if not the same as
+            model_name.
+        processor (Optional[str]): Pretrained processor name or path if not the same as
+            model_name.
+        cache_dir (Optional[str]): Where to store the pretrained data from
+            huggingface.co.
+        use_auth_token (bool): Whether to use Hugging Face auth token for private
+            models.
         precision (str): Precision to cast model weights to, default to auto.
         tie_word_embeddings (bool): Whether the model's input and output word embeddings
             should be tied.
-        trust_remote_code_model (bool): Whether to allow for custom models to execute their
-            own modeling files.
+        trust_remote_code_model (bool): Whether to allow for custom models to execute
+            their own modeling files.
         save_compressed (bool): Whether to compress sparse models during save.
         oneshot_device (str): Device to run oneshot calibration on.
-        model_revision (str): The specific model version to use (can be branch name, tag, or commit id).
-        
+        model_revision (str): The specific model version to use (can be branch name,
+            tag, or commit id).
+
         # Recipe arguments
         recipe (Optional[str]): Path to a LLM Compressor sparsification recipe.
-        recipe_args (Optional[List[str]]): List of recipe arguments to evaluate, in the format
-            "key1=value1", "key2=value2".
-        clear_sparse_session (bool): Whether to clear CompressionSession/CompressionLifecycle
-            data between runs.
+        recipe_args (Optional[List[str]]): List of recipe arguments to evaluate, in the
+            format "key1=value1", "key2=value2".
+        clear_sparse_session (bool): Whether to clear CompressionSession/
+            CompressionLifecycle data between runs.
         stage (Optional[str]): The stage of the recipe to use for oneshot.
-        
+
         # Dataset arguments
-        dataset (Optional[str]): The name of the dataset to use (via the datasets library).
-        dataset_config_name (Optional[str]): The configuration name of the dataset to use.
+        dataset (Optional[str]): The name of the dataset to use (via the datasets
+            library).
+        dataset_config_name (Optional[str]): The configuration name of the dataset
+            to use.
         dataset_path (Optional[str]): Path to a custom dataset. Supports json, csv, dvc.
-        num_calibration_samples (int): Number of samples to use for one-shot calibration.
-        shuffle_calibration_samples (bool): Whether to shuffle the dataset before calibration.
+        num_calibration_samples (int): Number of samples to use for one-shot
+            calibration.
+        shuffle_calibration_samples (bool): Whether to shuffle the dataset before
+            calibration.
         max_seq_length (int): Maximum total input sequence length after tokenization.
         pad_to_max_length (bool): Whether to pad all samples to `max_seq_length`.
         text_column (str): Key to use as the `text` input to tokenizer/processor.
-        concatenate_data (bool): Whether to concatenate datapoints to fill max_seq_length.
+        concatenate_data (bool): Whether to concatenate datapoints to fill
+            max_seq_length.
         streaming (bool): True to stream data from a cloud dataset.
         overwrite_cache (bool): Whether to overwrite the cached preprocessed datasets.
-        preprocessing_num_workers (Optional[int]): Number of processes for preprocessing.
-        min_tokens_per_module (Optional[float]): Minimum percentage of tokens per module,
-            relevant for MoE models.
-        trust_remote_code_data (bool): Whether to allow for datasets defined on the Hub using
-            a dataset script.
-        
+        preprocessing_num_workers (Optional[int]): Number of processes for
+            preprocessing.
+        min_tokens_per_module (Optional[float]): Minimum percentage of tokens per
+            module, relevant for MoE models.
+        trust_remote_code_data (bool): Whether to allow for datasets defined on the Hub
+            using a dataset script.
+
         # Output arguments
         output_dir (Optional[str]): Path to save the output model after calibration.
-    
+
     Returns:
         PreTrainedModel: The calibrated model
     """
@@ -289,13 +297,11 @@ def oneshot(
         "save_compressed": save_compressed,
         "oneshot_device": oneshot_device,
         "model_revision": model_revision,
-        
         # Recipe parameters
         "recipe": recipe,
         "recipe_args": recipe_args,
         "clear_sparse_session": clear_sparse_session,
         "stage": stage,
-        
         # Dataset parameters
         "dataset": dataset,
         "dataset_config_name": dataset_config_name,
@@ -311,18 +317,17 @@ def oneshot(
         "preprocessing_num_workers": preprocessing_num_workers,
         "min_tokens_per_module": min_tokens_per_module,
         "trust_remote_code_data": trust_remote_code_data,
-        
         # Output parameters
         "output_dir": output_dir,
     }
-    
+
     params = {k: v for k, v in params.items() if v is not None}
-    
+
     # Merge with any kwargs (this preserves backward compatibility)
     # kwargs take precedence over explicit params if same key exists
     all_args = {**params, **kwargs}
-    
+
     one_shot = Oneshot(**all_args)
     one_shot()
-    
+
     return one_shot.model