Relax dependency on accelerate and datasets in OVQuantizer

Author: eaidova

optimum/intel/openvino/quantization.py

@@ -22,8 +22,6 @@
 import openvino
 import torch
 import transformers
-from accelerate.data_loader import DataLoaderStateMixin
-from datasets import Dataset, load_dataset
 from nncf import NNCFConfig, compress_weights
 from nncf.torch import create_compressed_model, register_default_init_args, register_module
 from nncf.torch.dynamic_graph.io_handling import wrap_nncf_model_inputs_with_objwalk
@@ -33,13 +31,15 @@
 from torch.utils.data import DataLoader, RandomSampler
 from transformers import DataCollator, PreTrainedModel, default_data_collator
 from transformers.pytorch_utils import Conv1D
+from transformers.utils import is_accelerate_available
 
 from optimum.exporters.tasks import TasksManager
 from optimum.quantization_base import OptimumQuantizer
 
 from ...exporters.openvino import export, export_pytorch_via_onnx
 from ...exporters.openvino.stateful import ensure_export_task_support_stateful
 from ..utils.constant import _TASK_ALIASES
+from ..utils.import_utils import DATASETS_IMPORT_ERROR, is_datasets_available
 from .configuration import OVConfig
 from .modeling_base import OVBaseModel
 from .modeling_decoder import OVBaseDecoderModel
@@ -51,6 +51,9 @@
 )
 
 
+if is_datasets_available():
+    from datasets import Dataset
+
 COMPRESSION_OPTIONS = {
     "int8": {"mode": nncf.CompressWeightsMode.INT8},
     "int4_sym_g128": {"mode": nncf.CompressWeightsMode.INT4_SYM, "group_size": 128},
@@ -72,8 +75,11 @@ def get_inputs(self, dataloader_output) -> Tuple[Tuple, Dict]:
     @property
     def batch_size(self):
         batch_size = self._data_loader.batch_size
-        if batch_size is None and isinstance(self._data_loader, DataLoaderStateMixin):
-            batch_size = self._data_loader.total_batch_size
+        if is_accelerate_available():
+            from accelerate.data_loader import DataLoaderStateMixin
+
+            if batch_size is None and isinstance(self._data_loader, DataLoaderStateMixin):
+                batch_size = self._data_loader.total_batch_size
         return batch_size
 
 
@@ -155,7 +161,7 @@ def from_pretrained(cls, model: PreTrainedModel, **kwargs):
 
     def quantize(
         self,
-        calibration_dataset: Dataset = None,
+        calibration_dataset: "Dataset" = None,
         save_directory: Union[str, Path] = None,
         quantization_config: OVConfig = None,
         file_name: Optional[str] = None,
@@ -268,7 +274,7 @@ def _get_compression_options(self, config: OVConfig):
 
     def _quantize_ovbasemodel(
         self,
-        calibration_dataset: Dataset,
+        calibration_dataset: "Dataset",
         save_directory: Union[str, Path],
         batch_size: int = 1,
         data_collator: Optional[DataCollator] = None,
@@ -304,7 +310,7 @@ def _quantize_ovbasemodel(
 
     def _quantize_ovcausallm(
         self,
-        calibration_dataset: Dataset,
+        calibration_dataset: "Dataset",
         save_directory: Union[str, Path],
         batch_size: int = 1,
         data_collator: Optional[DataCollator] = None,
@@ -358,7 +364,7 @@ def _quantize_ovcausallm(
 
     def _quantize_torchmodel(
         self,
-        calibration_dataset: Dataset,
+        calibration_dataset: "Dataset",
         save_directory: Union[str, Path],
         quantization_config: OVConfig = None,
         file_name: Optional[str] = None,
@@ -482,7 +488,7 @@ def get_calibration_dataset(
         preprocess_batch: bool = True,
         use_auth_token: bool = False,
         cache_dir: Optional[str] = None,
-    ) -> Dataset:
+    ) -> "Dataset":
         """
         Create the calibration `datasets.Dataset` to use for the post-training static quantization calibration step.
 
@@ -507,6 +513,10 @@ def get_calibration_dataset(
         Returns:
             The calibration `datasets.Dataset` to use for the post-training static quantization calibration step.
         """
+        if not is_datasets_available():
+            raise ValueError(DATASETS_IMPORT_ERROR.format("OVQuantizer.get_calibration_dataset"))
+        from datasets import load_dataset
+
         calibration_dataset = load_dataset(
             dataset_name,
             name=dataset_config_name,
@@ -526,7 +536,7 @@ def get_calibration_dataset(
 
     def _get_calibration_dataloader(
         self,
-        calibration_dataset: Dataset,
+        calibration_dataset: "Dataset",
         batch_size: int,
         remove_unused_columns: bool,
         data_collator: Optional[DataCollator] = None,
@@ -543,6 +553,6 @@ def _get_calibration_dataloader(
         )
         return OVDataLoader(calibration_dataloader)
 
-    def _remove_unused_columns(self, dataset: Dataset):
+    def _remove_unused_columns(self, dataset: "Dataset"):
         ignored_columns = list(set(dataset.column_names) - set(self._signature_columns))
         return dataset.remove_columns(ignored_columns)

optimum/intel/utils/import_utils.py

@@ -119,6 +119,16 @@
         _timm_available = False
 
 
+_datasets_available = importlib.util.find_spec("datasets") is not None
+_datasets_version = "N/A"
+
+if _datasets_available:
+    try:
+        _datasets_version = importlib_metadata.version("datasets")
+    except importlib_metadata.PackageNotFoundError:
+        _datasets_available = False
+
+
 def is_transformers_available():
     return _transformers_available
 
@@ -151,6 +161,10 @@ def is_timm_available():
     return _timm_available
 
 
+def is_datasets_available():
+    return _datasets_available
+
+
 # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319
 def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
     """
@@ -267,6 +281,11 @@ def is_timm_version(operation: str, version: str):
 `pip install neural-compressor`. Please note that you may need to restart your runtime after installation.
 """
 
+DATASETS_IMPORT_ERROR = """
+{0} requires the datasets library but it was not found in your environment. You can install it with pip:
+`pip install datasets`. Please note that you may need to restart your runtime after installation.
+"""
+
 BACKENDS_MAPPING = OrderedDict(
     [
         ("diffusers", (is_diffusers_available, DIFFUSERS_IMPORT_ERROR)),

setup.py

@@ -45,7 +45,7 @@
         "transformers>=4.34.0",
     ],
     "openvino": ["openvino>=2023.2", "onnx", "onnxruntime", "transformers>=4.36.0", "optimum>=1.16.1"],
-    "nncf": ["nncf>=2.7.0"],
+    "nncf": ["nncf>=2.7.0", "datasets", "accelerate"],
     "ipex": ["intel-extension-for-pytorch", "onnx"],
     "diffusers": ["diffusers"],
     "quality": QUALITY_REQUIRE,