Refactor through OVQuantizer call

Author: nikita-savelyevv

optimum/intel/openvino/modeling_decoder.py

@@ -625,21 +625,19 @@ def _from_pretrained(
                     "Quantization of the weights requires nncf, please install it with `pip install nncf`"
                 )
 
-            from optimum.intel.openvino.quantization import _weight_only_quantization
+            from optimum.intel.openvino.quantization import OVQuantizer
 
             default_config = _check_default_4bit_configs(config)
+
             if default_config:
                 logger.info(
                     f"For the given model, we recommend the following `quantization_config` : {default_config}"
                 )
 
+            quantizer = OVQuantizer(causal_model)
             quantization_config_copy = copy.deepcopy(quantization_config)
             quantization_config_copy.tokenizer = quantization_config.tokenizer or model_id
-            _weight_only_quantization(
-                model,
-                quantization_config_copy,
-                transform_fn=lambda x: causal_model.prepare_inputs(**x),
-            )
+            quantizer.quantize(ov_config=OVConfig(quantization_config=quantization_config_copy))
 
         return causal_model
 

optimum/intel/openvino/quantization.py

@@ -198,7 +198,7 @@ def from_pretrained(cls, model: PreTrainedModel, **kwargs):
     def quantize(
         self,
         calibration_dataset: Optional[Union[datasets.Dataset, nncf.Dataset, Iterable]] = None,
-        save_directory: Union[str, Path] = None,
+        save_directory: Optional[Union[str, Path]] = None,
         ov_config: OVConfig = None,
         file_name: Optional[str] = None,
         batch_size: int = 1,
@@ -214,7 +214,7 @@ def quantize(
             calibration_dataset (`datasets.Dataset` or `nncf.Dataset` or `Iterable`, *optional*):
                 A collection of data samples to use for quantization calibration. Is optional for weight-only
                 quantization and is required for full quantization.
-            save_directory (`Union[str, Path]`):
+            save_directory (`Union[str, Path]`, *optional*):
                 The directory where the quantized model should be saved.
             ov_config (`OVConfig`, *optional*):
                 The configuration containing the parameters related to quantization. If not provided, 8-bit symmetric
@@ -262,10 +262,6 @@ def quantize(
                 "as an instance of `OVWeightQuantizationConfig` for weight-only compression or as an instance of `OVQuantizationConfig` for full model quantization."
             )
 
-        if save_directory is None:
-            # TODO : can be set to self.model.config.name_or_path for OVModels when not provided
-            raise ValueError("`save_directory` needs to be specified")
-
         if ov_config is None:
             ov_config = OVConfig()
         if not isinstance(ov_config, OVConfig):
@@ -318,21 +314,41 @@ def quantize(
     def _quantize_ovbasemodel(
         self,
         ov_config: OVConfig,
-        save_directory: Union[str, Path],
+        save_directory: Union[str, Path] = None,
         calibration_dataset: Optional[Union[datasets.Dataset, nncf.Dataset, Iterable]] = None,
         batch_size: int = 1,
         data_collator: Optional[DataCollator] = None,
         remove_unused_columns: bool = True,
         **kwargs,
     ):
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
+        if save_directory is not None:
+            save_directory = Path(save_directory)
+            save_directory.mkdir(parents=True, exist_ok=True)
 
         quantization_config = ov_config.quantization_config
         if isinstance(quantization_config, OVWeightQuantizationConfig):
+            if calibration_dataset is None and isinstance(quantization_config.dataset, str):
+                from optimum.intel import OVModelForCausalLM
+
+                if isinstance(self.model, OVModelForCausalLM):
+                    from optimum.gptq.data import get_dataset, prepare_dataset
+
+                    tokenizer = AutoTokenizer.from_pretrained(quantization_config.tokenizer)
+                    nsamples = quantization_config.num_samples if quantization_config.num_samples else 128
+                    calibration_dataset = get_dataset(
+                        quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples
+                    )
+                    calibration_dataset = prepare_dataset(calibration_dataset)
+                    calibration_dataset = nncf.Dataset(calibration_dataset, lambda x: self.model.prepare_inputs(**x))
+                else:
+                    raise ValueError(
+                        f"Can't create weight compression calibration dataset from string for {type(self.model)}"
+                    )
+
             _weight_only_quantization(self.model.model, quantization_config, calibration_dataset)
-            self.model.save_pretrained(save_directory)
-            ov_config.save_pretrained(save_directory)
+            if save_directory is not None:
+                self.model.save_pretrained(save_directory)
+                ov_config.save_pretrained(save_directory)
             return
         if not isinstance(quantization_config, OVQuantizationConfig):
             raise ValueError(f"Unsupported type of quantization config: {type(quantization_config)}")
@@ -384,8 +400,9 @@ def _quantize_ovbasemodel(
             **kwargs,
         )
         self.model.model = quantized_model
-        self.model.save_pretrained(save_directory)
-        ov_config.save_pretrained(save_directory)
+        if save_directory is not None:
+            self.model.save_pretrained(save_directory)
+            ov_config.save_pretrained(save_directory)
 
     def _quantize_torchmodel(
         self,
@@ -398,6 +415,10 @@ def _quantize_torchmodel(
         remove_unused_columns: bool = True,
         **kwargs,
     ):
+        if save_directory is None:
+            # TODO : can be set to self.model.config.name_or_path for OVModels when not provided
+            raise ValueError("`save_directory` needs to be specified")
+
         self._set_task()
         save_directory = Path(save_directory)
         save_directory.mkdir(parents=True, exist_ok=True)
@@ -622,7 +643,6 @@ def _weight_only_quantization(
     model: openvino.runtime.Model,
     quantization_config: Union[OVWeightQuantizationConfig, Dict],
     calibration_dataset: Optional[Union[nncf.Dataset, Iterable]] = None,
-    transform_fn: Optional[Callable] = None,
 ) -> openvino.runtime.Model:
     config = quantization_config
     if isinstance(config, dict):
@@ -645,15 +665,6 @@ def _weight_only_quantization(
             dataset = calibration_dataset
         else:
             dataset = nncf.Dataset(calibration_dataset)
-    elif config.dataset is not None and isinstance(config.dataset, str):
-        tokenizer = AutoTokenizer.from_pretrained(config.tokenizer)
-
-        from optimum.gptq.data import get_dataset, prepare_dataset
-
-        nsamples = config.num_samples if config.num_samples else 128
-        dataset = get_dataset(config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
-        dataset = prepare_dataset(dataset)
-        dataset = nncf.Dataset(dataset, transform_fn)
 
     sensitivity_metric = None
     if isinstance(config.sensitivity_metric, str):