From 7e95327edf283d5b0d501b240e107dac1efe0b14 Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Tue, 17 Oct 2023 16:53:16 +0200 Subject: [PATCH] format --- optimum/intel/neural_compressor/modeling_base.py | 6 ++++-- optimum/intel/neural_compressor/modeling_decoder.py | 10 ++++++++-- tests/neural_compressor/test_modeling.py | 8 ++++---- 3 files changed, 16 insertions(+), 8 deletions(-) diff --git a/optimum/intel/neural_compressor/modeling_base.py b/optimum/intel/neural_compressor/modeling_base.py index cc3953007c..8746a8d817 100644 --- 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 index 80851b1f43..e284ce4c3e 100644 --- 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 index 93e849d04a..854ae61ac0 100644 --- 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: