apply style

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

examples/quantization_w4a16/llama3_example.py

@@ -75,4 +75,4 @@ def tokenize(sample):
 # Save to disk compressed.
 SAVE_DIR = model_id.split("/")[-1] + "-W4A16-G128"
 model.save_pretrained(SAVE_DIR, save_compressed=True)
-tokenizer.save_pretrained(SAVE_DIR)
+tokenizer.save_pretrained(SAVE_DIR)

src/llmcompressor/utils/dev.py

@@ -2,20 +2,24 @@
 import logging
 import os
 import tempfile
-from typing import Type, Dict, Any, Union
+from typing import Type
 
 import torch
+from accelerate import dispatch_model, infer_auto_device_map
+from accelerate.utils import get_balanced_memory
 from huggingface_hub import snapshot_download
 from safetensors.torch import save_file
 from transformers import AutoModelForCausalLM, PreTrainedModel
 from transformers.modeling_utils import TORCH_INIT_FUNCTIONS
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, WEIGHTS_INDEX_NAME
-from accelerate import dispatch_model, infer_auto_device_map
-from accelerate.utils import get_balanced_memory
 
 from llmcompressor.utils.helpers import patch_attr
 
-__all__ = ["skip_weights_download", "patch_transformers_logger_level", "dispatch_for_generation"]
+__all__ = [
+    "skip_weights_download",
+    "patch_transformers_logger_level",
+    "dispatch_for_generation",
+]
 
 
 @contextlib.contextmanager
@@ -114,8 +118,8 @@ def dispatch_for_generation(model: PreTrainedModel) -> PreTrainedModel:
     max_memory = get_balanced_memory(
         model,
         dtype=model.dtype,
-        no_split_module_classes=model._get_no_split_modules("auto")
+        no_split_module_classes=model._get_no_split_modules("auto"),
     )
     device_map = infer_auto_device_map(model, dtype=model.dtype, max_memory=max_memory)
 
-    return dispatch_model(model, device_map=device_map)
+    return dispatch_model(model, device_map=device_map)