Fix quantization config dataset constraint (#694)

huggingface · Apr 30, 2024 · b017856 · b017856
1 parent c235ae1
commit b017856
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/optimum/intel/openvino/configuration.py b/optimum/intel/openvino/configuration.py
@@ -133,7 +133,7 @@ class OVWeightQuantizationConfig(OVQuantizationConfigBase):
                     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 ['wikitext','c4','c4-new','ptb','ptb-new'] for LLLMs
+            in a list of strings or just use the one from the list ['wikitext2','c4','c4-new','ptb','ptb-new'] for LLLMs
             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.
@@ -194,7 +194,7 @@ def post_init(self):
                 f"If you wish to provide a custom dataset, please use the `OVQuantizer` instead."
             )
         if self.dataset is not None and isinstance(self.dataset, str):
-            llm_datasets = ["wikitext", "c4", "c4-new", "ptb", "ptb-new"]
+            llm_datasets = ["wikitext2", "c4", "c4-new", "ptb", "ptb-new"]
             stable_diffusion_datasets = [
                 "conceptual_captions",
                 "laion/220k-GPT4Vision-captions-from-LIVIS",

diff --git a/tests/openvino/test_quantization.py b/tests/openvino/test_quantization.py
@@ -715,7 +715,7 @@ class OVQuantizationConfigTest(unittest.TestCase):
         (OVWeightQuantizationConfig(bits=8, sym=True),),
         (
             OVWeightQuantizationConfig(
-                dataset="wikitext",
+                dataset="wikitext2",
                 bits=4,
                 ignored_scope={"names": ["op_name"]},
                 sym=False,
@@ -747,7 +747,7 @@ class OVQuantizationConfigTest(unittest.TestCase):
         (dict(bits=8, sym=True), OVWeightQuantizationConfig, None),
         (
             dict(
-                dataset="wikitext",
+                dataset="wikitext2",
                 bits=4,
                 ignored_scope={"names": ["op_name"]},
                 sym=False,
@@ -771,7 +771,7 @@ class OVQuantizationConfigTest(unittest.TestCase):
         (dict(num_samples=100), OVWeightQuantizationConfig, "Can't determine type of OV quantization config"),
         (dict(abc="def"), OVWeightQuantizationConfig, "Can't determine type of OV quantization config"),
         (
-            dict(bits=8, fast_bias_correction=True, dataset="wikitext"),
+            dict(bits=8, fast_bias_correction=True, dataset="wikitext2"),
             OVWeightQuantizationConfig,
             "Can't determine type of OV quantization config",
         ),
@@ -793,7 +793,7 @@ class OVQuantizationConfigTest(unittest.TestCase):
         (dict(abc="def", weight_only=False), OVQuantizationConfig, None),
         (dict(abc="def", weight_only=True), OVWeightQuantizationConfig, None),
         (
-            dict(bits=8, fast_bias_correction=True, dataset="wikitext", weight_only=True),
+            dict(bits=8, fast_bias_correction=True, dataset="wikitext2", weight_only=True),
             OVWeightQuantizationConfig,
             None,
         ),