Implement dataset preparation from list of strings

Author: nikita-savelyevv

optimum/intel/openvino/configuration.py

@@ -221,11 +221,11 @@ class OVWeightQuantizationConfig(OVQuantizationConfigBase):
                 - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
                     using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
         dataset (`str or List[str]`, *optional*):
-            The dataset used for data-aware compression or quantization with NNCF. You can provide your own dataset
-            in a list of strings or just use the one from the list ['wikitext2','c4','c4-new'] for language models
-            or ['conceptual_captions','laion/220k-GPT4Vision-captions-from-LIVIS','laion/filtered-wit'] for diffusion models.
-            Alternatively, you can provide data objects via `calibration_dataset` argument
-            of `OVQuantizer.quantize()` method.
+            The dataset used for data-aware compression with NNCF. For language models you can provide your own dataset
+            in a list of strings or just use the one from the list ['wikitext2','c4','c4-new']. For diffusion models it
+            must be one of ['conceptual_captions', 'laion/220k-GPT4Vision-captions-from-LIVIS', 'laion/filtered-wit'].
+            Alternatively, you can provide data objects via `calibration_dataset` argument of `OVQuantizer.quantize()`
+            method.
         ratio (`float`, defaults to 1.0):
             The ratio between baseline and backup precisions (e.g. 0.9 means 90% of layers quantized to INT4_ASYM
             and the rest to INT8_ASYM).

optimum/intel/openvino/quantization.py

@@ -351,7 +351,7 @@ def _quantize_ovbasemodel(
                     "quantization. Will rely on `calibration_dataset`."
                 )
 
-            if calibration_dataset is None and isinstance(quantization_config.dataset, str):
+            if calibration_dataset is None and quantization_config.dataset is not None:
                 from optimum.intel import OVModelForCausalLM
 
                 if isinstance(self.model, OVModelForCausalLM):
@@ -676,7 +676,12 @@ def _prepare_builtin_dataset(self, quantization_config: OVWeightQuantizationConf
             quantization_config.tokenizer, trust_remote_code=quantization_config.trust_remote_code
         )
         nsamples = quantization_config.num_samples if quantization_config.num_samples else 128
-        calibration_dataset = get_dataset(quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
+        if isinstance(quantization_config.dataset, str):
+            calibration_dataset = get_dataset(quantization_config.dataset, tokenizer, seqlen=32, nsamples=nsamples)
+        else:
+            calibration_dataset = [
+                tokenizer(text, return_tensors="pt") for text in quantization_config.dataset[:nsamples]
+            ]
         calibration_dataset = prepare_dataset(calibration_dataset)
         calibration_dataset = nncf.Dataset(calibration_dataset, lambda x: self.model.prepare_inputs(**x))
 

tests/openvino/test_quantization.py

@@ -223,6 +223,19 @@ class OVWeightCompressionTest(unittest.TestCase):
             ),
             14,
         ),
+        (
+            OVModelForCausalLM,
+            "opt",
+            dict(
+                bits=4,
+                sym=True,
+                group_size=-1,
+                ratio=0.8,
+                sensitivity_metric="mean_activation_magnitude",
+                dataset=["one two, " * i for i in range(130)],
+            ),
+            14,
+        ),
         (
             OVModelForCausalLM,
             "llama_awq",