Make tests alive

Author: nikita-savelyevv

optimum/exporters/openvino/__main__.py

@@ -501,7 +501,7 @@ class StoreAttr(object):
                     "Quantization of the weights requires nncf, please install it with `pip install nncf`"
                 )
 
-            from optimum.intel.openvino.quantization.ov_quantizer import _weight_only_quantization
+            from optimum.intel.openvino.quantization.quantizer import _weight_only_quantization
 
             _weight_only_quantization(submodel, quantization_config)
             compressed_submodel_path = submodel_path.parent / f"{submodel_path.stem}_compressed.xml"

optimum/intel/openvino/modeling_base.py

@@ -259,7 +259,7 @@ def fix_op_names_duplicates(model: openvino.runtime.Model):
                     "Quantization of the weights to int8 requires nncf, please install it with `pip install nncf`"
                 )
 
-            from optimum.intel.openvino.quantization.ov_quantizer import _weight_only_quantization
+            from optimum.intel.openvino.quantization.quantizer import _weight_only_quantization
 
             if not isinstance(quantization_config, (dict, OVWeightQuantizationConfig)):
                 raise TypeError(

optimum/intel/openvino/modeling_decoder.py

@@ -891,7 +891,7 @@ def _from_pretrained(
                     "quantization is not supported with `compile_only` mode, please intialize model without this option"
                 )
 
-            from optimum.intel.openvino.quantization.ov_quantizer import OVQuantizer
+            from optimum.intel.openvino.quantization.quantizer import OVQuantizer
 
             default_config = _check_default_4bit_configs(config.name_or_path)
 

optimum/intel/openvino/modeling_visual_language.py

@@ -561,7 +561,7 @@ def _from_pretrained(
         )
 
         if to_quantize:
-            from optimum.intel.openvino.quantization.ov_quantizer import OVQuantizer
+            from optimum.intel.openvino.quantization.quantizer import OVQuantizer
 
             quantization_config_copy = copy.deepcopy(quantization_config)
             quantization_config_copy.tokenizer = quantization_config.tokenizer or model_id

optimum/intel/openvino/quantization/__init__.py

@@ -19,6 +19,7 @@
     OVMixedQuantizationConfig,
     OVQuantizationConfig,
     OVQuantizationConfigBase,
+    OVQuantizationMethod,
     OVWeightQuantizationConfig,
 )
-from .ov_quantizer import OVQuantizer
+from .quantizer import OVQuantizer

optimum/intel/openvino/quantization/dataset_builder.py optimum/intel/openvino/quantization/calibration_dataset_builder.py

@@ -25,10 +25,10 @@
 
 from optimum.configuration_utils import BaseConfig
 
-from ..utils.import_utils import is_nncf_available
-from .utils import (
+from ...utils.import_utils import is_nncf_available
+from ..utils import (
     LANGUAGE_DATASETS,
-    PREDEFINED_SD_DATASETS,
+    PREDEFINED_DIFFUSION_DATASETS,
     PREDEFINED_SPEECH_TO_TEXT_DATASETS,
     PREDEFINED_VISUAL_LM_DATASETS,
 )
@@ -514,7 +514,7 @@ def post_init(self):
             )
         if self.dataset is not None and isinstance(self.dataset, str):
             visual_lm_datasets = list(PREDEFINED_VISUAL_LM_DATASETS.keys())
-            stable_diffusion_datasets = list(PREDEFINED_SD_DATASETS.keys())
+            stable_diffusion_datasets = list(PREDEFINED_DIFFUSION_DATASETS.keys())
             if self.dataset not in LANGUAGE_DATASETS + visual_lm_datasets + stable_diffusion_datasets:
                 raise ValueError(
                     f"""You have entered a string value for dataset. You can only choose between
@@ -523,7 +523,7 @@ def post_init(self):
                 )
 
         if self.dataset is not None and not (
-            self.quant_method == OVQuantizationMethod.AWQ
+            self.quant_method in [OVQuantizationMethod.AWQ, OVQuantizationMethod.HYBRID]
             or self.scale_estimation
             or self.gptq
             or self.lora_correction

optimum/intel/openvino/quantization/configuration.py

@@ -51,22 +51,19 @@
     ONNX_WEIGHTS_NAME,
     OV_XML_FILE_NAME,
 )
+from .calibration_dataset_builder import OVCalibrationDatasetBuilder
 from .configuration import (
     OVConfig,
     OVMixedQuantizationConfig,
     OVQuantizationConfig,
     OVQuantizationMethod,
     OVWeightQuantizationConfig,
 )
-from .dataset_builder import OVCalibrationDatasetBuilder
 
 
 if is_datasets_available():
     from datasets import Dataset
 
-if is_diffusers_available():
-    from optimum.intel.openvino.modeling_diffusion import OVDiffusionPipeline
-
 register_module(ignored_algorithms=[])(Conv1D)
 
 core = Core()
@@ -177,6 +174,9 @@ def quantize(
                 "quantization. Will rely on `calibration_dataset`."
             )
 
+        if is_diffusers_available():
+            from optimum.intel.openvino.modeling_diffusion import OVDiffusionPipeline
+
         if calibration_dataset is not None and isinstance(calibration_dataset, nncf.Dataset):
             # TODO: add a warning that calibration dataset should be a dictionary in this case
             calibration_dataset = {"model": calibration_dataset}
@@ -192,6 +192,14 @@ def quantize(
                 # TODO: analyze this execution path
                 calibration_dataset = calibration_dataset.select_columns(["caption"])
 
+            if (
+                isinstance(self.model, OVDiffusionPipeline)
+                and isinstance(calibration_dataset, list)
+                and all(isinstance(it, str) for it in calibration_dataset)
+            ):
+                # TODO: deprecate this way of providing calibration dataset
+                data_collator = data_collator or (lambda x: x)
+
             calibration_dataset = self.dataset_builder.build_from_dataset(
                 quantization_config, calibration_dataset, batch_size, data_collator, remove_unused_columns
             )
@@ -233,6 +241,9 @@ def _quantize_ovbasemodel(
         from optimum.intel.openvino.modeling_seq2seq import _OVModelForWhisper
         from optimum.intel.openvino.modeling_visual_language import OVModelForVisualCausalLM
 
+        if is_diffusers_available():
+            from optimum.intel.openvino.modeling_diffusion import OVDiffusionPipeline
+
         quantization_config = ov_config.quantization_config
         if calibration_datasets is None and quantization_config.dataset is not None:
             calibration_datasets = self.dataset_builder.build_from_quantization_config(quantization_config)
@@ -290,7 +301,7 @@ def _quantize_ovbasemodel(
                         raise ValueError("Calibration datasets for Diffusion models should contain only one value.")
                     # Apply hybrid quantization to diffusion model
                     diffusion_model_name, calibration_dataset = next(iter(calibration_datasets.items()))
-                    diffusion_model = getattr(self.model, diffusion_model_name)
+                    diffusion_model = getattr(self.model, diffusion_model_name).model
                     getattr(self.model, diffusion_model_name).model = _hybrid_quantization(
                         diffusion_model, quantization_config, calibration_dataset, **kwargs
                     )
@@ -476,7 +487,7 @@ def _set_task(self):
     def get_calibration_dataset(
         self,
         dataset_name: str,
-        num_samples: int = 100,
+        num_samples: Optional[int] = None,
         dataset_config_name: Optional[str] = None,
         dataset_split: str = "train",
         preprocess_function: Optional[Callable] = None,
@@ -492,7 +503,7 @@ def get_calibration_dataset(
             dataset_name (`str`):
                 The dataset repository name on the Hugging Face Hub or path to a local directory containing data files
                 in generic formats and optionally a dataset script, if it requires some code to read the data files.
-            num_samples (`int`, defaults to 100):
+            num_samples (`int`, *optional*):
                 The maximum number of samples composing the calibration dataset.
             dataset_config_name (`str`, *optional*):
                 The name of the dataset configuration.
@@ -518,6 +529,7 @@ def get_calibration_dataset(
 
         return self.dataset_builder.load_dataset(
             dataset_name,
+            num_samples,
             dataset_config_name,
             dataset_split,
             preprocess_function,
@@ -531,11 +543,7 @@ def get_calibration_dataset(
 def _quantize_whisper_model(
     model, quantization_config: OVQuantizationConfig, calibration_dataset: Dict[str, nncf.Dataset], **kwargs
 ):
-    for submodel_name in ["encoder_model", "decoder_model", "decoder_with_past_model"]:
-        submodel = getattr(model, submodel_name, None)
-        if submodel_name == "decoder_with_past_model" and submodel is None:
-            continue
-
+    for submodel_name, submodel in model.ov_submodels.items():
         config = quantization_config.clone()
         # quantization_config.num_samples of audio samples result in more actual model inputs
         config.num_samples = calibration_dataset[submodel_name].get_length()

optimum/intel/openvino/quantization/ov_quantizer.py optimum/intel/openvino/quantization/quantizer.py

@@ -141,13 +141,13 @@
 LANGUAGE_DATASETS = ["wikitext2", "c4", "c4-new", "auto"]
 
 PREDEFINED_DIFFUSION_DATASETS = {
-    "conceptual_captions": {"split": "train", "inputs": {"prompt": "caption"}, "streaming": True},
+    "conceptual_captions": {"split": "train", "prompt_column_name": "caption", "streaming": True},
     "laion/220k-GPT4Vision-captions-from-LIVIS": {
         "split": "train",
-        "inputs": {"prompt": "caption"},
+        "prompt_column_name": "caption",
         "streaming": True,
     },
-    "laion/filtered-wit": {"split": "train", "inputs": {"prompt": "caption"}, "streaming": True},
+    "laion/filtered-wit": {"split": "train", "prompt_column_name": "caption", "streaming": True},
 }
 
 PREDEFINED_VISUAL_LM_DATASETS = {
@@ -163,7 +163,6 @@
         "id": "openslr/librispeech_asr",
         "name": "clean",
         "split": "validation",
-        "inputs": {"audio": ("audio", "array"), "sampling_rate": ("audio", "sampling_rate")},
         "streaming": True,
     }
 }

optimum/intel/openvino/utils.py

@@ -78,7 +78,7 @@
 from optimum.intel.openvino.utils import TemporaryDirectory
 from copy import deepcopy
 
-from optimum.intel.openvino.quantization.ov_quantizer import InferRequestWrapper
+from optimum.intel.openvino.quantization.calibration_dataset_builder import InferRequestWrapper
 from optimum.intel.utils.import_utils import is_openvino_version, is_transformers_version
 from utils_tests import (
     MODEL_NAMES,
@@ -1398,6 +1398,8 @@ def test_config_init_kwargs(
         ]
     )
     def test_quantization_kwargs_override(self, mock_method_name, quantization_function, dataset_key, config_type):
+        from optimum.intel.openvino.quantization.quantizer import _weight_only_quantization, _full_quantization
+
         with unittest.mock.patch(mock_method_name) as mock_method:
             mock_model = unittest.mock.Mock([])
             mock_model.get_rt_info = unittest.mock.Mock(return_value={})
@@ -1407,7 +1409,11 @@ def test_quantization_kwargs_override(self, mock_method_name, quantization_funct
 
             additional_kwargs = {"param2": "new_value2", "param3": "value3"}
 
-            quantization_function = globals()[quantization_function]
+            quantization_function = (
+                _weight_only_quantization
+                if quantization_function == "_weight_only_quantization"
+                else _full_quantization
+            )
             quantization_function(mock_model, mock_quantization_config, None, **additional_kwargs)
 
             expected_kwargs = {"param1": "value1", "param2": "new_value2", "param3": "value3", dataset_key: None}