Add support for single npz file with multiple samples

Signed-off-by: gcunhase <4861122+gcunhase@users.noreply.github.com>

Author: gcunhase

modelopt/onnx/autocast/referencerunner.py

@@ -30,6 +30,7 @@
 import onnx
 
 from modelopt.onnx.autocast.logging_config import configure_logging, logger
+from modelopt.onnx.quantization.calib_utils import CalibrationDataProvider
 from modelopt.onnx.quantization.ort_utils import _prepare_ep_list
 
 configure_logging()
@@ -70,7 +71,11 @@ def _load_inputs_from_json(self, input_data_path):
 
     def _load_inputs_from_npz(self, input_data_path):
         """Load inputs from NPZ format."""
-        return [np.load(input_data_path)]
+        calib_data = np.load(input_data_path)
+
+        # Wrap data into a CalibDataProvider to support a single NPZ file containing data from multiple batches
+        data_loader = {key: calib_data[key] for key in calib_data.files}
+        return CalibrationDataProvider(self.model, data_loader).calibration_data_list
 
     def _validate_inputs(self, data_loader):
         """Validate that input names and shapes match the model."""

modelopt/onnx/quantization/calib_utils.py

@@ -38,7 +38,7 @@ class CalibrationDataProvider(CalibrationDataReader):
 
     def __init__(
         self,
-        onnx_path: str,
+        onnx_path: str | onnx.ModelProto,
         calibration_data: CalibrationDataType,
         calibration_shapes: str | None = None,
     ):
@@ -58,7 +58,7 @@ def __init__(
         logger.info("Setting up CalibrationDataProvider for calibration")
         # Tensor data is not required to generate the calibration data
         # So even if the model has external data, we don't need to load them here
-        onnx_model = onnx.load(onnx_path)
+        onnx_model = onnx.load(onnx_path) if isinstance(onnx_path, str) else onnx_path
         input_names = get_input_names(onnx_model)
         input_shapes = {} if calibration_shapes is None else parse_shapes_spec(calibration_shapes)
         inferred_input_shapes = get_input_shapes(onnx_model)