format

huggingface · Oct 17, 2023 · 7e95327 · 7e95327
1 parent 15c89c2
commit 7e95327
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/optimum/intel/neural_compressor/modeling_base.py b/optimum/intel/neural_compressor/modeling_base.py
@@ -79,7 +79,9 @@ def __init__(
         self.inc_config = inc_config
         self._q_config = q_config
         self.model_save_dir = model_save_dir
-        self._device = getattr(self.model, "device", None) or torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        self._device = getattr(self.model, "device", None) or torch.device(
+            "cuda:0" if torch.cuda.is_available() else "cpu"
+        )
 
         if getattr(self.config, "backend", None) == "ipex":
             if not is_ipex_available():
@@ -176,7 +178,7 @@ def _from_pretrained(
             model, config=config, model_save_dir=model_save_dir, q_config=q_config, inc_config=inc_config, **kwargs
         )
 
-    def _save_pretrained(self, save_directory: Union[str, Path], file_name : str = WEIGHTS_NAME):
+    def _save_pretrained(self, save_directory: Union[str, Path], file_name: str = WEIGHTS_NAME):
         output_path = os.path.join(save_directory, file_name)
 
         if isinstance(self.model, torch.nn.Module):

diff --git a/optimum/intel/neural_compressor/modeling_decoder.py b/optimum/intel/neural_compressor/modeling_decoder.py
@@ -15,7 +15,7 @@
 import logging
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Optional, Union, Dict
+from typing import Dict, Optional, Union
 
 from transformers import AutoModelForCausalLM, PretrainedConfig
 from transformers.file_utils import add_start_docstrings
@@ -53,5 +53,11 @@ def __init__(
         **kwargs,
     ):
         super(INCModelForCausalLM, self).__init__(
-            model=model, config=config, model_save_dir=model_save_dir, q_config=q_config, inc_config=inc_config, use_cache=use_cache, **kwargs
+            model=model,
+            config=config,
+            model_save_dir=model_save_dir,
+            q_config=q_config,
+            inc_config=inc_config,
+            use_cache=use_cache,
+            **kwargs,
         )
diff --git a/tests/neural_compressor/test_modeling.py b/tests/neural_compressor/test_modeling.py
@@ -15,8 +15,8 @@
 
 import os
 import tempfile
-import unittest
 import time
+import unittest
 
 import torch
 from parameterized import parameterized
@@ -66,7 +66,6 @@
 DIFFUSERS_MODEL_NAMES_TO_TASK = (("echarlaix/stable-diffusion-v1-5-inc-int8-dynamic", "stable-diffusion"),)
 
 
-
 class Timer(object):
     def __enter__(self):
         self.elapsed = time.perf_counter()
@@ -76,7 +75,6 @@ def __exit__(self, type, value, traceback):
         self.elapsed = (time.perf_counter() - self.elapsed) * 1e3
 
 
-
 class INCModelingTest(unittest.TestCase):
     GENERATION_LENGTH = 100
     SPEEDUP_CACHE = 1.1
@@ -148,7 +146,9 @@ def test_compare_with_and_without_past_key_values(self):
             outputs_model_with_pkv = model_with_pkv.generate(
                 **tokens, min_length=self.GENERATION_LENGTH, max_length=self.GENERATION_LENGTH, num_beams=1
             )
-        model_without_pkv = INCModelForCausalLM.from_pretrained(model_id, use_cache=False, subfolder="model_without_pkv")
+        model_without_pkv = INCModelForCausalLM.from_pretrained(
+            model_id, use_cache=False, subfolder="model_without_pkv"
+        )
         # Warmup
         model_without_pkv.generate(**tokens)
         with Timer() as without_pkv_timer: