style

docker/cpu.dockerfile

@@ -1,6 +1,5 @@
 FROM ubuntu:latest
 
-
 # Ignore interactive questions during `docker build`
 ENV DEBIAN_FRONTEND noninteractive
 

examples/api_launch.py

@@ -1,8 +1,9 @@
-from optimum_benchmark.logging_utils import setup_logging
-from optimum_benchmark.experiment import launch, ExperimentConfig
 from optimum_benchmark.backends.pytorch.config import PyTorchConfig
-from optimum_benchmark.launchers.torchrun.config import TorchrunConfig
 from optimum_benchmark.benchmarks.inference.config import InferenceConfig
+from optimum_benchmark.experiment import ExperimentConfig, launch
+from optimum_benchmark.launchers.torchrun.config import TorchrunConfig
+from optimum_benchmark.logging_utils import setup_logging
+
 
 if __name__ == "__main__":
     setup_logging(level="INFO")

optimum_benchmark/backends/base.py

@@ -1,25 +1,25 @@
 import gc
 import random
 from abc import ABC
-from logging import getLogger
 from collections import OrderedDict
-from typing import Optional, ClassVar, Generic, Dict, Any
+from logging import getLogger
+from typing import Any, ClassVar, Dict, Generic, Optional
+
+import numpy as np
+from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel, TrainerState
 
-from .config import BackendConfigT
 from ..task_utils import get_automodel_class_for_task
+from .config import BackendConfigT
 from .diffusers_utils import extract_diffusers_shapes_from_config, get_diffusers_pretrained_config
-from .timm_utils import extract_timm_shapes_from_config, get_timm_pretrained_config, get_timm_pre_processor
+from .timm_utils import extract_timm_shapes_from_config, get_timm_pre_processor, get_timm_pretrained_config
 from .transformers_utils import (
+    PretrainedProcessor,
     extract_transformers_shapes_from_artifacts,
     get_transformers_generation_config,
-    get_transformers_pretrained_config,
     get_transformers_pre_processor,
-    PretrainedProcessor,
+    get_transformers_pretrained_config,
 )
 
-import numpy as np
-from transformers import GenerationConfig, PretrainedConfig, PreTrainedModel, TrainerState
-
 LOGGER = getLogger("backend")
 
 

optimum_benchmark/backends/config.py

@@ -1,14 +1,14 @@
 import os
 from abc import ABC
-from logging import getLogger
 from dataclasses import dataclass, field
-from typing import Optional, TypeVar, Dict, Any
+from logging import getLogger
+from typing import Any, Dict, Optional, TypeVar
+
+from psutil import cpu_count
 
 from ..system_utils import get_gpu_device_ids, is_nvidia_system, is_rocm_system
 from ..task_utils import infer_library_from_model_name_or_path, infer_task_from_model_name_or_path
 
-from psutil import cpu_count
-
 LOGGER = getLogger("backend")
 
 HUB_KWARGS = {"revision": "main", "force_download": False, "local_files_only": False, "trust_remote_code": False}

optimum_benchmark/backends/diffusers_utils.py

@@ -1,9 +1,9 @@
 from typing import Dict
 
-from ..import_utils import is_diffusers_available
-
 from hydra.utils import get_class
 
+from ..import_utils import is_diffusers_available
+
 if is_diffusers_available():
     import diffusers  # type: ignore
 

optimum_benchmark/backends/neural_compressor/backend.py

@@ -1,22 +1,22 @@
-import os
 import gc
-from typing import Any, Dict
+import os
 from logging import getLogger
 from tempfile import TemporaryDirectory
-
-from ...generators.dataset_generator import DatasetGenerator
-from ..transformers_utils import randomize_weights
-from .utils import TASKS_TO_INCMODELS
-from .config import INCConfig
-from ..base import Backend
+from typing import Any, Dict
 
 import torch
 from hydra.utils import get_class
-from transformers.utils import ModelOutput
+from neural_compressor.config import AccuracyCriterion, PostTrainingQuantConfig, TuningCriterion
+from optimum.intel.neural_compressor.quantization import INCQuantizer
 from transformers.modeling_utils import no_init_weights
+from transformers.utils import ModelOutput
 from transformers.utils.logging import set_verbosity_error
-from optimum.intel.neural_compressor.quantization import INCQuantizer
-from neural_compressor.config import PostTrainingQuantConfig, AccuracyCriterion, TuningCriterion
+
+from ...generators.dataset_generator import DatasetGenerator
+from ..base import Backend
+from ..transformers_utils import randomize_weights
+from .config import INCConfig
+from .utils import TASKS_TO_INCMODELS
 
 # disable transformers logging
 set_verbosity_error()

optimum_benchmark/backends/neural_compressor/config.py

@@ -1,10 +1,10 @@
-from typing import Any, Dict, Optional
 from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
 
 from omegaconf import OmegaConf
 
-from ..config import BackendConfig
 from ...import_utils import neural_compressor_version
+from ..config import BackendConfig
 
 # https://github.com/intel/neural-compressor/blob/master/neural_compressor/config.py#L490
 ACCURACY_CRITERION_CONFIG = {"higher_is_better": True, "criterion": "relative", "tolerable_loss": 0.01}

optimum_benchmark/backends/onnxruntime/backend.py

@@ -1,40 +1,40 @@
 import gc
 import os
-from logging import getLogger
 from collections import OrderedDict
+from logging import getLogger
 from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List
 
-from ..base import Backend
-from .config import ORTConfig
-from ...task_utils import TEXT_GENERATION_TASKS
-from ...generators.dataset_generator import DatasetGenerator
-from .utils import format_calibration_config, format_quantization_config, TASKS_TO_ORTMODELS, TASKS_TO_ORTSD
-
 import torch
 from datasets import Dataset
 from hydra.utils import get_class
 from onnxruntime import SessionOptions
-from safetensors.torch import save_file
-from transformers import TrainerCallback
-from transformers.modeling_utils import no_init_weights
-from transformers.utils.logging import set_verbosity_error
-from optimum.onnxruntime.configuration import (
-    AutoOptimizationConfig,
-    AutoQuantizationConfig,
-    AutoCalibrationConfig,
-    OptimizationConfig,
-    QuantizationConfig,
-    CalibrationConfig,
-)
 from optimum.onnxruntime import (
-    ONNX_DECODER_WITH_PAST_NAME,
     ONNX_DECODER_NAME,
-    ORTTrainingArguments,
+    ONNX_DECODER_WITH_PAST_NAME,
     ORTOptimizer,
     ORTQuantizer,
     ORTTrainer,
+    ORTTrainingArguments,
+)
+from optimum.onnxruntime.configuration import (
+    AutoCalibrationConfig,
+    AutoOptimizationConfig,
+    AutoQuantizationConfig,
+    CalibrationConfig,
+    OptimizationConfig,
+    QuantizationConfig,
 )
+from safetensors.torch import save_file
+from transformers import TrainerCallback
+from transformers.modeling_utils import no_init_weights
+from transformers.utils.logging import set_verbosity_error
+
+from ...generators.dataset_generator import DatasetGenerator
+from ...task_utils import TEXT_GENERATION_TASKS
+from ..base import Backend
+from .config import ORTConfig
+from .utils import TASKS_TO_ORTMODELS, TASKS_TO_ORTSD, format_calibration_config, format_quantization_config
 
 # disable transformers logging
 set_verbosity_error()

optimum_benchmark/backends/onnxruntime/config.py

@@ -1,9 +1,9 @@
 import os
-from typing import Any, Dict, Optional
 from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
 
-from ..config import BackendConfig
 from ...import_utils import onnxruntime_version
+from ..config import BackendConfig
 from ..peft_utils import PEFT_CONFIGS, PEFT_TASKS_TYPES
 
 QUANTIZATION_CONFIG = {

optimum_benchmark/backends/onnxruntime/utils.py

@@ -1,8 +1,7 @@
 from typing import Any, Dict
 
+from onnxruntime.quantization import CalibrationMethod, QuantFormat, QuantizationMode, QuantType
 from optimum.pipelines import ORT_SUPPORTED_TASKS
-from onnxruntime.quantization import CalibrationMethod, QuantizationMode, QuantFormat, QuantType
-
 
 TASKS_TO_ORTSD = {
     "stable-diffusion": "optimum.onnxruntime.ORTStableDiffusionPipeline",

optimum_benchmark/backends/openvino/backend.py

@@ -1,26 +1,26 @@
 import gc
-import os
 import inspect
-from typing import Any, Dict
-from logging import getLogger
+import os
 from collections import OrderedDict
+from logging import getLogger
 from tempfile import TemporaryDirectory
-
-from ..base import Backend
-from .config import OVConfig
-from .utils import TASKS_TO_OVMODEL
-from ...task_utils import TEXT_GENERATION_TASKS
-from ..transformers_utils import randomize_weights
-from ...generators.dataset_generator import DatasetGenerator
+from typing import Any, Dict
 
 import torch
 from hydra.utils import get_class
 from openvino.runtime import properties
-from safetensors.torch import save_file
+from optimum.intel.openvino import OVConfig as OVQuantizationConfig  # naming conflict
 from optimum.intel.openvino import OVQuantizer
+from safetensors.torch import save_file
 from transformers.modeling_utils import no_init_weights
 from transformers.utils.logging import set_verbosity_error
-from optimum.intel.openvino import OVConfig as OVQuantizationConfig  # naming conflict
+
+from ...generators.dataset_generator import DatasetGenerator
+from ...task_utils import TEXT_GENERATION_TASKS
+from ..base import Backend
+from ..transformers_utils import randomize_weights
+from .config import OVConfig
+from .utils import TASKS_TO_OVMODEL
 
 # disable transformers logging
 set_verbosity_error()

optimum_benchmark/backends/openvino/config.py

@@ -1,8 +1,8 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
-from ..config import BackendConfig
 from ...import_utils import openvino_version
+from ..config import BackendConfig
 
 
 @dataclass

optimum_benchmark/backends/peft_utils.py

@@ -4,10 +4,10 @@
 
 if is_peft_available():
     from peft import (
+        AdaLoraConfig,
         IA3Config,
         LoraConfig,
         PeftConfig,
-        AdaLoraConfig,
         PrefixTuningConfig,
         PromptEncoderConfig,
         PromptLearningConfig,

optimum_benchmark/backends/pytorch/backend.py

@@ -1,23 +1,23 @@
 import gc
 import os
-from logging import getLogger
 from collections import OrderedDict
+from logging import getLogger
 from tempfile import TemporaryDirectory
 from typing import Any, Callable, Dict, List
 
-from ..base import Backend
-from .config import PyTorchConfig
-from ..peft_utils import get_peft_config_class
-from ..transformers_utils import randomize_weights
-from ...import_utils import is_deepspeed_available, is_peft_available, is_torch_distributed_available
-
+import datasets.utils.logging as datasets_logging
 import torch
+import transformers.utils.logging as transformers_logging
 from datasets import Dataset
 from safetensors.torch import save_file
-import datasets.utils.logging as datasets_logging
+from transformers import Trainer, TrainerCallback, TrainerState, TrainingArguments
 from transformers.modeling_utils import no_init_weights
-import transformers.utils.logging as transformers_logging
-from transformers import TrainerCallback, TrainerState, Trainer, TrainingArguments
+
+from ...import_utils import is_deepspeed_available, is_peft_available, is_torch_distributed_available
+from ..base import Backend
+from ..peft_utils import get_peft_config_class
+from ..transformers_utils import randomize_weights
+from .config import PyTorchConfig
 
 if is_peft_available():
     from peft import get_peft_model  # type: ignore

optimum_benchmark/backends/pytorch/config.py

@@ -1,9 +1,9 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, Optional
 
-from ..config import BackendConfig
-from ...system_utils import is_rocm_system
 from ...import_utils import torch_version
+from ...system_utils import is_rocm_system
+from ..config import BackendConfig
 from ..peft_utils import PEFT_CONFIGS, PEFT_TASKS_TYPES
 
 DEVICE_MAPS = ["auto", "sequential"]

optimum_benchmark/backends/tensorrt_llm/backend.py

@@ -1,13 +1,13 @@
 from logging import getLogger
 from typing import Any, Dict
 
+from hydra.utils import get_class
+from transformers.utils import ModelOutput
+
 from ..base import Backend
 from .config import TRTLLMConfig
 from .utils import MODEL_TYPE_TO_TRTLLMMODEL
 
-from hydra.utils import get_class
-from transformers.utils import ModelOutput
-
 LOGGER = getLogger("tensorrt-llm")
 
 

optimum_benchmark/backends/tensorrt_llm/config.py

@@ -1,9 +1,8 @@
-from typing import Optional
 from dataclasses import dataclass
+from typing import Optional
 
-from ..config import BackendConfig
 from ...import_utils import tesnorrt_llm_version
-
+from ..config import BackendConfig
 
 SUPPORTED_DTYPES = ["float16", "bfloat16", "float32"]
 

optimum_benchmark/backends/text_generation_inference/backend.py

@@ -1,23 +1,24 @@
 import gc
 import os
 import time
+from concurrent.futures import ThreadPoolExecutor
 from logging import getLogger
-from typing import Any, Dict, List
 from tempfile import TemporaryDirectory
-from concurrent.futures import ThreadPoolExecutor
-
-from ..base import Backend
-from .config import TGIConfig
-from ...task_utils import TEXT_GENERATION_TASKS
-from ..transformers_utils import randomize_weights
+from typing import Any, Dict, List
 
 import torch
-import docker
-import docker.types
-import docker.errors
-from safetensors.torch import save_model
 from huggingface_hub import InferenceClient, snapshot_download
 from huggingface_hub.inference._text_generation import TextGenerationResponse
+from safetensors.torch import save_model
+
+import docker
+import docker.errors
+import docker.types
+
+from ...task_utils import TEXT_GENERATION_TASKS
+from ..base import Backend
+from ..transformers_utils import randomize_weights
+from .config import TGIConfig
 
 # bachend logger
 LOGGER = getLogger("text-generation-inference")

optimum_benchmark/backends/timm_utils.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional
 
-from ..import_utils import is_timm_available, is_transformers_available, is_torch_available
+from ..import_utils import is_timm_available, is_torch_available, is_transformers_available
 
 if is_torch_available():
     import torch