diff --git a/optimum_benchmark/backends/ipex/backend.py b/optimum_benchmark/backends/ipex/backend.py index d32a203d..a637ced0 100644 --- a/optimum_benchmark/backends/ipex/backend.py +++ b/optimum_benchmark/backends/ipex/backend.py @@ -9,7 +9,7 @@ from ..base import Backend from ..transformers_utils import fast_weights_init from .config import IPEXConfig -from .utils import TASKS_TO_IPEXMODEL +from .utils import TASKS_TO_IPEXMODELS if is_accelerate_available(): from accelerate import Accelerator @@ -24,8 +24,8 @@ class IPEXBackend(Backend[IPEXConfig]): def __init__(self, config: IPEXConfig) -> None: super().__init__(config) - if self.config.task in TASKS_TO_IPEXMODEL: - self.ipexmodel_class = get_class(TASKS_TO_IPEXMODEL[self.config.task]) + if self.config.task in TASKS_TO_IPEXMODELS: + self.ipexmodel_class = get_class(TASKS_TO_IPEXMODELS[self.config.task]) self.logger.info(f"\t+ Using IPEXModel class {self.ipexmodel_class.__name__}") else: raise NotImplementedError(f"IPEXBackend does not support task {self.config.task}") diff --git a/optimum_benchmark/backends/ipex/utils.py b/optimum_benchmark/backends/ipex/utils.py index dd68428e..7b4a83f8 100644 --- a/optimum_benchmark/backends/ipex/utils.py +++ b/optimum_benchmark/backends/ipex/utils.py @@ -1,4 +1,4 @@ -TASKS_TO_IPEXMODEL = { +TASKS_TO_IPEXMODELS = { "fill-mask": "optimum.intel.IPEXModelForMaskedLM", "text-generation": "optimum.intel.IPEXModelForCausalLM", "feature-extraction": "optimum.intel.IPEXModel", diff --git a/optimum_benchmark/backends/onnxruntime/backend.py b/optimum_benchmark/backends/onnxruntime/backend.py index b50e2258..ce7386b3 100644 --- a/optimum_benchmark/backends/onnxruntime/backend.py +++ b/optimum_benchmark/backends/onnxruntime/backend.py @@ -28,7 +28,7 @@ from .config import ORTConfig from .utils import ( TASKS_TO_ORTMODELS, - TASKS_TO_ORTPIPELINE, + TASKS_TO_ORTPIPELINES, format_calibration_config, format_quantization_config, ) @@ -49,8 +49,8 @@ def __init__(self, config: ORTConfig) -> None: if self.config.library != "diffusers" and self.config.task in TASKS_TO_ORTMODELS: self.ort_model_loader = get_class(TASKS_TO_ORTMODELS[self.config.task]) self.logger.info(f"Using ORTModel class {self.ort_model_loader.__name__}") - elif self.config.library == "diffusers" and self.config.task in TASKS_TO_ORTPIPELINE: - self.ort_model_loader = get_class(TASKS_TO_ORTPIPELINE[self.config.task]) + elif self.config.library == "diffusers" and self.config.task in TASKS_TO_ORTPIPELINES: + self.ort_model_loader = get_class(TASKS_TO_ORTPIPELINES[self.config.task]) self.logger.info(f"Using ORTDiffusionPipeline class {self.ort_model_loader.__name__}") else: raise NotImplementedError(f"ORTBackend does not support task {self.config.task}") diff --git a/optimum_benchmark/backends/onnxruntime/utils.py b/optimum_benchmark/backends/onnxruntime/utils.py index 63598223..e8cbe1eb 100644 --- a/optimum_benchmark/backends/onnxruntime/utils.py +++ b/optimum_benchmark/backends/onnxruntime/utils.py @@ -7,7 +7,7 @@ task: f"optimum.onnxruntime.{task_dict['class'][0].__name__}" for task, task_dict in ORT_SUPPORTED_TASKS.items() } -TASKS_TO_ORTPIPELINE = { +TASKS_TO_ORTPIPELINES = { "inpainting": "optimum.onnxruntime.ORTPipelineForInpainting", "text-to-image": "optimum.onnxruntime.ORTPipelineForText2Image", "image-to-image": "optimum.onnxruntime.ORTPipelineForImage2Image", diff --git a/optimum_benchmark/backends/openvino/backend.py b/optimum_benchmark/backends/openvino/backend.py index be903e7d..11be6e13 100644 --- a/optimum_benchmark/backends/openvino/backend.py +++ b/optimum_benchmark/backends/openvino/backend.py @@ -9,7 +9,7 @@ from ..base import Backend from ..transformers_utils import fast_weights_init from .config import OVConfig as OVBackendConfig -from .utils import TASKS_OVPIPELINE, TASKS_TO_OVMODEL +from .utils import TASKS_TO_OVMODELS, TASKS_TO_OVPIPELINES if is_accelerate_available(): from accelerate import Accelerator @@ -24,11 +24,11 @@ class OVBackend(Backend[OVBackendConfig]): def __init__(self, config: OVBackendConfig) -> None: super().__init__(config) - if self.config.task in TASKS_TO_OVMODEL: - self.ovmodel_class = get_class(TASKS_TO_OVMODEL[self.config.task]) + if self.config.library != "diffusers" and self.config.task in TASKS_TO_OVMODELS: + self.ovmodel_class = get_class(TASKS_TO_OVMODELS[self.config.task]) self.logger.info(f"\t+ Using OVModel class {self.ovmodel_class.__name__}") - elif self.config.task in TASKS_OVPIPELINE: - self.ovmodel_class = get_class(TASKS_OVPIPELINE[self.config.task]) + elif self.config.library == "diffusers" and self.config.task in TASKS_TO_OVPIPELINES: + self.ovmodel_class = get_class(TASKS_TO_OVPIPELINES[self.config.task]) self.logger.info(f"\t+ Using OVDiffusionPipeline class {self.ovmodel_class.__name__}") else: raise NotImplementedError(f"OVBackend does not support task {self.config.task}") diff --git a/optimum_benchmark/backends/openvino/utils.py b/optimum_benchmark/backends/openvino/utils.py index 51f9629c..e382d724 100644 --- a/optimum_benchmark/backends/openvino/utils.py +++ b/optimum_benchmark/backends/openvino/utils.py @@ -1,4 +1,4 @@ -TASKS_TO_OVMODEL = { +TASKS_TO_OVMODELS = { "fill-mask": "optimum.intel.openvino.OVModelForMaskedLM", "text-generation": "optimum.intel.openvino.OVModelForCausalLM", "text2text-generation": "optimum.intel.openvino.OVModelForSeq2SeqLM", @@ -10,7 +10,7 @@ "audio-classification": "optimum.intel.openvino.OVModelForAudioClassification", "pix2struct": "optimum.intel.openvino.OVModelForPix2Struct", } -TASKS_OVPIPELINE = { +TASKS_TO_OVPIPELINES = { "inpainting": "optimum.intel.openvino.OVPipelineForInpainting", "text-to-image": "optimum.intel.openvino.OVPipelineForText2Image", "image-to-image": "optimum.intel.openvino.OVPipelineForImage2Image", diff --git a/optimum_benchmark/backends/tensorrt_llm/backend.py b/optimum_benchmark/backends/tensorrt_llm/backend.py index 5e7bad89..8cd046eb 100644 --- a/optimum_benchmark/backends/tensorrt_llm/backend.py +++ b/optimum_benchmark/backends/tensorrt_llm/backend.py @@ -12,7 +12,7 @@ from ..base import Backend from ..transformers_utils import fast_weights_init from .config import TRTLLMConfig -from .utils import MODEL_TYPE_TO_TRTLLMMODEL +from .utils import MODEL_TYPE_TO_TRTLLMMODELS class TRTLLMBackend(Backend[TRTLLMConfig]): @@ -21,8 +21,8 @@ class TRTLLMBackend(Backend[TRTLLMConfig]): def __init__(self, config: TRTLLMConfig): super().__init__(config) - if self.config.model_type in MODEL_TYPE_TO_TRTLLMMODEL: - self.trtllm_loader = get_class(MODEL_TYPE_TO_TRTLLMMODEL[self.config.model_type]) + if self.config.model_type in MODEL_TYPE_TO_TRTLLMMODELS: + self.trtllm_loader = get_class(MODEL_TYPE_TO_TRTLLMMODELS[self.config.model_type]) self.logger.info(f"\t+ Using TRTLLMModel class {self.trtllm_loader.__name__}") else: raise NotImplementedError(f"TRTLLMBackend does not support model_type {self.config.model_type}") diff --git a/optimum_benchmark/backends/tensorrt_llm/utils.py b/optimum_benchmark/backends/tensorrt_llm/utils.py index 4574da53..01b6ed0e 100644 --- a/optimum_benchmark/backends/tensorrt_llm/utils.py +++ b/optimum_benchmark/backends/tensorrt_llm/utils.py @@ -1 +1 @@ -MODEL_TYPE_TO_TRTLLMMODEL = {"llama": "optimum.nvidia.models.llama.LlamaForCausalLM"} +MODEL_TYPE_TO_TRTLLMMODELS = {"llama": "optimum.nvidia.models.llama.LlamaForCausalLM"}