add datafiles

huggingface · Oct 3, 2023 · 8a2874b · 8a2874b
1 parent 8ee658d
commit 8a2874b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/optimum/onnxruntime/quantization.py b/optimum/onnxruntime/quantization.py
@@ -482,6 +482,7 @@ def quantize(
     def get_calibration_dataset(
         self,
         dataset_name: str,
+        data_files: str = None,
         num_samples: int = 100,
         dataset_config_name: Optional[str] = None,
         dataset_split: Optional[str] = None,
@@ -525,6 +526,7 @@ def get_calibration_dataset(
         calib_dataset = load_dataset(
             dataset_name,
             name=dataset_config_name,
+            data_files=data_files,
             split=dataset_split,
             use_auth_token=use_auth_token,
         )