ai-edge-quantizer-nightly 0.0.1.dev20250115__py3-none-any.whl

ai_edge_quantizer/calibrator.py

@@ -0,0 +1,288 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Quantization Calibration."""
+
+from collections.abc import Callable, Iterable
+import copy
+from typing import Any, Union
+
+from absl import logging
+import numpy as np
+
+from ai_edge_quantizer import algorithm_manager
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer import recipe_manager
+from ai_edge_quantizer.utils import calibration_utils
+from ai_edge_quantizer.utils import tfl_flatbuffer_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+_SignatureInput = dict[str, Any]  # input_argument_name -> tensor_value.
+_SignatureOutput = dict[
+    str, np.ndarray
+]  # output_argument_name -> tensor_value.
+
+
+class Calibrator:
+  """Calibrator for TFLite model."""
+
+  def __init__(
+      self,
+      float_tflite: Union[str, bytes],
+      num_threads: int = 16,
+  ):
+    self._flatbuffer_model = tfl_flatbuffer_utils.read_model(float_tflite)
+
+    if not tfl_flatbuffer_utils.is_float_model(self._flatbuffer_model):
+      raise ValueError(
+          "The input model for calibration is not a float model. Please check"
+          " the model (e.g., if it is already quantized)."
+      )
+    self._tfl_interpreter = tfl_interpreter_utils.create_tfl_interpreter(
+        float_tflite, use_xnnpack=True, num_threads=num_threads
+    )
+    # Tensor name to tensor content.
+    self._tensor_content_map: dict[str, Any] = {}
+    # QSV of all the tensors in the model.
+    self._model_qsvs: dict[str, qtyping.QSV] = {}
+    # Cached output of the model.
+    self._cached_output: list[_SignatureOutput] = []
+
+  # TODO(b/330740605)- Collect multiple QSVs in one run to save compute.
+  def calibrate(
+      self,
+      calibration_dataset: dict[str, Iterable[_SignatureInput]],
+      model_recipe_manager: recipe_manager.RecipeManager,
+      cache_output: bool = False,
+      qsv_update_func: Callable[
+          [qtyping.QSV, qtyping.QSV],
+          qtyping.QSV,
+      ] = calibration_utils.moving_average_update,
+  ) -> None:
+    """Calibrates the model using the given dataset for a model signature.
+
+    The process is
+    0. Initialize quantization statistics values (QSVs) using the initialization
+    function (from AlgorithmManager) for the op if needed.
+    1. Invoke TFL interpreter on the calibration data.
+    2. Go through each op, ask RecipeManager for the quantization setting
+    of the op.
+    3. Ask AlgorithmManager for the calibration function of the op given the
+    quantization setting.
+    4. Apply the function to the op to obtain quantization statistics (qsvs) for
+    the tensors associated with the op.
+    5. Update the global qsv dictionary
+    6. Start another round of calibration.
+
+    Args:
+      calibration_dataset: A dictionary of input data for calibration for the
+        given model signature.
+      model_recipe_manager: A RecipeManager object that contains the
+        quantization recipe.
+      cache_output: Whether to cache the output of the model during the
+        calibration process. This is useful if there are dependencies between
+        signatures/models (e.g., decode requires encode output).
+      qsv_update_func: The function to update the QSVs.
+    """
+    op_codes = self._flatbuffer_model.operatorCodes
+    if not self._model_qsvs:
+      self._initialize_model_qsvs(model_recipe_manager)
+    else:
+      logging.warning(
+          "Calibrator contains non-empty model qsvs, and the current"
+          " calibration process will start on top of this state (i.e., update"
+          " the existing qsvs). If this is an unintended behavior please call"
+          " reset_model_qsvs to reset model qsvs."
+      )
+
+    # TODO: b/329322226 - Enable parallel calibration.
+    for signature_key, dataset in calibration_dataset.items():
+      # Step0: get subgraph index.
+      subgraph_idx = tfl_interpreter_utils.get_signature_main_subgraph_index(
+          self._tfl_interpreter, signature_key
+      )
+
+      for data in dataset:
+        # Initialize tensor names that are updated in this round of calibration.
+        updated_tensor_names = set()
+
+        # Step1: run tfl interpreter on subgraph to get tensor content.
+        signature_output = tfl_interpreter_utils.invoke_interpreter_signature(
+            self._tfl_interpreter, data, signature_key
+        )
+        if cache_output:
+          self._cached_output.append(signature_output)
+        self._tensor_content_map.update(
+            tfl_interpreter_utils.get_tensor_name_to_content_map(
+                self._tfl_interpreter, subgraph_idx
+            )
+        )
+
+        # Step2: go through each op in subgraph to update quantization
+        # statistic values.
+        subgraph = self._flatbuffer_model.subgraphs[subgraph_idx]
+        graph_info = qtyping.GraphInfo(
+            subgraph.tensors, self._flatbuffer_model.buffers
+        )
+        # Add input/output operators to the subgraph.
+        subgraph.operators += (
+            tfl_flatbuffer_utils.get_subgraph_input_output_operators(subgraph)
+        )
+        for op in subgraph.operators:
+          if isinstance(op, qtyping.IOOperator):
+            op_key = op.op_key
+          else:
+            op_code = op_codes[op.opcodeIndex].builtinCode
+            if op_code not in tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME:
+              continue
+            op_key = tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME[op_code]
+          # Step2.1: query the quantization_recipe to get op quantization
+          # settings.
+          op_scope = self._get_op_scope(op, subgraph.tensors)
+          algorithm_name, _ = model_recipe_manager.get_quantization_configs(
+              op_key, op_scope
+          )
+          if algorithm_name == algorithm_manager.AlgorithmName.NO_QUANTIZE:
+            continue
+          # Step2.2: query algorithm_manager to get/call the related calibration
+          # function.
+          calibrate_func = algorithm_manager.get_quantization_func(
+              algorithm_name, op_key, qtyping.QuantizeMode.CALIBRATE
+          )
+          op_qsvs = calibrate_func(op, graph_info, self._tensor_content_map)
+          # Step3: Update tensor qsvs with the new values. Ignore the tensor
+          # names that are already updated in this round of calibration.
+          op_updated_tensor_name = self._update_qsvs(
+              op_qsvs, updated_tensor_names, qsv_update_func
+          )
+          updated_tensor_names.update(op_updated_tensor_name)
+      # Reset interpreter after one round of calibration.
+      self._tfl_interpreter.reset_all_variables()
+
+  def get_model_qsvs(self) -> dict[str, qtyping.QSV]:
+    """Get the model qsvs.
+
+    Returns:
+      A dictionary of tensor name to QSV.
+    """
+    return self._model_qsvs
+
+  def get_cached_output(self) -> list[_SignatureOutput]:
+    """Get the cached output of the model."""
+    return self._cached_output
+
+  def clear_cached_output(self) -> None:
+    """Clear the cached output of the model."""
+    self._cached_output = []
+
+  def reset_model_qsvs(self) -> None:
+    """Reset the model qsvs."""
+    self._model_qsvs = {}
+
+  def load_model_qsvs(self, model_qsvs: dict[str, qtyping.QSV]) -> None:
+    """Load the model qsvs.
+
+    Args:
+      model_qsvs: A dictionary of tensor name to QSV.
+    """
+    self._model_qsvs = copy.deepcopy(model_qsvs)
+
+  def _update_qsvs(
+      self,
+      op_qsvs: dict[str, qtyping.QSV],
+      ignore_tensor_names: set[str],
+      qsv_update_func: Callable[[qtyping.QSV, qtyping.QSV], qtyping.QSV],
+  ) -> set[str]:
+    """Update the model qsvs with the new values.
+
+    Args:
+      op_qsvs: A dictionary of tensor name to QSV.
+      ignore_tensor_names: A set of tensor names to ignore.
+      qsv_update_func: The function to update the QSVs.
+
+    Returns:
+      A set of tensor names that are updated.
+    """
+    updated_tensor_names = set()
+    for tensor_name, qsv in op_qsvs.items():
+      if tensor_name in ignore_tensor_names:
+        continue
+      if tensor_name not in self._model_qsvs:
+        self._model_qsvs[tensor_name] = qsv
+      else:
+        updated_qsv = qsv_update_func(self._model_qsvs[tensor_name], qsv)
+        self._model_qsvs[tensor_name] = updated_qsv
+      updated_tensor_names.add(tensor_name)
+    return updated_tensor_names
+
+  def _get_op_scope(self, op, subgraph_tensors) -> str:
+    """Get the scope of the op.
+
+    The scope is the name of the output tensor of the op.
+
+    Args:
+      op: The op to get the scope.
+      subgraph_tensors: The tensors in the subgraph.
+
+    Returns:
+      The scope of the op.
+    """
+    scope = ""
+    for output_tensor_idx in op.outputs:
+      if output_tensor_idx != -1:
+        output_tensor = subgraph_tensors[output_tensor_idx]
+        scope += tfl_flatbuffer_utils.get_tensor_name(output_tensor)
+    return scope
+
+  # TODO: b/354224138 - Remove code duplication between calibrate and
+  # _initialize_model_qsvs.
+  def _initialize_model_qsvs(
+      self, model_recipe_manager: recipe_manager.RecipeManager
+  ) -> None:
+    """Initialize the model qsvs.
+
+    Args:
+      model_recipe_manager: A RecipeManager object that contains the
+        quantization recipe.
+    """
+    op_codes = self._flatbuffer_model.operatorCodes
+    for subgraph in self._flatbuffer_model.subgraphs:
+      graph_info = qtyping.GraphInfo(
+          subgraph.tensors, self._flatbuffer_model.buffers
+      )
+      for subgraph_op_id, op in enumerate(subgraph.operators):
+        op_code = op_codes[op.opcodeIndex].builtinCode
+        if op_code not in tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME:
+          continue
+        op_key = tfl_flatbuffer_utils.TFL_OP_CODE_TO_NAME[op_code]
+        # Step1: query the quantization_recipe to get op quantization
+        # settings.
+        op_scope = self._get_op_scope(op, subgraph.tensors)
+        algorithm_name, op_quant_config = (
+            model_recipe_manager.get_quantization_configs(op_key, op_scope)
+        )
+        if algorithm_name == algorithm_manager.AlgorithmName.NO_QUANTIZE:
+          continue
+        # Step2: query algorithm_manager to get/call the related qsv init
+        # function.
+        qsv_init_func = algorithm_manager.get_init_qsv_func(
+            algorithm_name, op_key
+        )
+        op_info = qtyping.OpInfo(op, op_key, subgraph_op_id, op_quant_config)
+        op_qsvs = qsv_init_func(op_info, graph_info)
+        # Step3: initialize tensor qsvs.
+        for tensor_name, qsv in op_qsvs.items():
+          if tensor_name not in self._model_qsvs:
+            self._model_qsvs[tensor_name] = qsv

ai_edge_quantizer/calibrator_test.py

@@ -0,0 +1,297 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for calibrator."""
+
+from collections.abc import Generator
+import os
+from typing import Any
+
+import numpy as np
+
+from tensorflow.python.platform import googletest
+from ai_edge_quantizer import calibrator
+from ai_edge_quantizer import qtyping
+from ai_edge_quantizer import recipe_manager
+from ai_edge_quantizer.utils import test_utils
+from ai_edge_quantizer.utils import tfl_interpreter_utils
+
+_ComputePrecision = qtyping.ComputePrecision
+_AlgorithmName = recipe_manager.AlgorithmName
+
+TEST_DATA_PREFIX_PATH = test_utils.get_path_to_datafile("")
+_TENSOR_QUANT_CONFIG = qtyping.TensorQuantizationConfig
+
+TEST_MIN_VAL, TEST_MAX_VAL = -1, 1
+
+_RNG = np.random.default_rng(66)
+
+
+def _representative_dataset_gen(size=(1, 8), num_samples=10):
+  for _ in range(num_samples):
+    vals = np.random.rand(*size).astype(np.float32)
+    vals[0][0], vals[0][1] = (
+        TEST_MIN_VAL,
+        TEST_MAX_VAL,
+    )  # fix min/max for testing
+    yield {"input_1": vals}
+
+
+def _get_calibration_data(
+    dataset_gen: Generator[dict[str, Any], Any, None],
+) -> dict[str, Any]:
+  calibration_samples = [sample for sample in dataset_gen]
+  calibration_data = {
+      tfl_interpreter_utils.DEFAULT_SIGNATURE_KEY: calibration_samples,
+  }
+  return calibration_data
+
+
+def _add_default_int8xint8_integer_recipe(recipe_manager_object):
+  recipe_manager_object.add_quantization_config(
+      regex=".*",
+      operation_name=qtyping.TFLOperationName.ALL_SUPPORTED,
+      algorithm_key=_AlgorithmName.MIN_MAX_UNIFORM_QUANT,
+      op_config=qtyping.OpQuantizationConfig(
+          activation_tensor_config=_TENSOR_QUANT_CONFIG(
+              num_bits=8, symmetric=False
+          ),
+          weight_tensor_config=_TENSOR_QUANT_CONFIG(num_bits=8, symmetric=True),
+          compute_precision=_ComputePrecision.INTEGER,
+      ),
+  )
+
+
+class CalibratorTest(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/single_fc.tflite"
+    )
+    self._calibrator = calibrator.Calibrator(self._test_model_path)
+    self._recipe_manager = recipe_manager.RecipeManager()
+    dataset_gen = _representative_dataset_gen()
+    self._representative_dataset = _get_calibration_data(dataset_gen)
+
+  def test_calibrator_state_manipulation(self):
+    # load/get qsvs
+    sample_qsv = {"serving_default_input_1:0": {"min": -10, "max": 8}}
+    self._calibrator.load_model_qsvs(sample_qsv)
+    model_tensor_qsvs = self._calibrator.get_model_qsvs()
+    self.assertLen(model_tensor_qsvs, 1)
+    self.assertIn("serving_default_input_1:0", model_tensor_qsvs)  # input
+    input_qsv = model_tensor_qsvs["serving_default_input_1:0"]
+    self.assertEqual(input_qsv["min"], -10)
+    self.assertEqual(input_qsv["max"], 8)
+
+    # reset qsvs
+    self._calibrator.reset_model_qsvs()
+    model_tensor_qsvs = self._calibrator.get_model_qsvs()
+    self.assertEmpty(model_tensor_qsvs)
+
+  def test_calibrator_initialize_qsv(self):
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    # Overwrite the single op to fc
+    self._recipe_manager.add_quantization_config(
+        regex=".*Stateful.*",
+        operation_name=qtyping.TFLOperationName.FULLY_CONNECTED,
+        algorithm_key=_AlgorithmName.MIN_MAX_UNIFORM_QUANT,
+        op_config=qtyping.OpQuantizationConfig(
+            weight_tensor_config=_TENSOR_QUANT_CONFIG(
+                num_bits=4,
+                granularity=qtyping.QuantGranularity.CHANNELWISE,
+            ),
+            compute_precision=_ComputePrecision.INTEGER,
+        ),
+    )
+    self._calibrator._initialize_model_qsvs(self._recipe_manager)
+    model_tensor_qsvs = self._calibrator.get_model_qsvs()
+
+    self.assertLen(model_tensor_qsvs, 4)
+    self.assertIn("serving_default_input_1:0", model_tensor_qsvs)  # input
+    input_qsv = model_tensor_qsvs["serving_default_input_1:0"]
+    self.assertEmpty(input_qsv)
+
+    self.assertIn("sequential/dense/MatMul", model_tensor_qsvs)  # weight
+    weight_tensor_qsv = model_tensor_qsvs["sequential/dense/MatMul"]
+    mins_maxs_shape = (16, 1)
+    self.assertTupleEqual(weight_tensor_qsv["min"].shape, mins_maxs_shape)
+    self.assertAlmostEqual(weight_tensor_qsv["min"][0][0], -0.40436327)
+    self.assertTupleEqual(weight_tensor_qsv["max"].shape, mins_maxs_shape)
+    self.assertAlmostEqual(weight_tensor_qsv["max"][0][0], 0.46138108)
+
+    self.assertIn(
+        "sequential/dense/BiasAdd/ReadVariableOp", model_tensor_qsvs
+    )  # bias
+    bias_tensor_qsv = model_tensor_qsvs[
+        "sequential/dense/BiasAdd/ReadVariableOp"
+    ]
+    mins_maxs_shape = (16,)
+    self.assertTupleEqual(bias_tensor_qsv["min"].shape, mins_maxs_shape)
+    self.assertAlmostEqual(bias_tensor_qsv["min"][0], -0.26978338)
+    self.assertTupleEqual(bias_tensor_qsv["max"].shape, mins_maxs_shape)
+    # Here bias min/max will be the same as each element is a scalar
+    # Bias will be quantized with input_scale * weight_scale.
+    self.assertSequenceEqual(
+        list(bias_tensor_qsv["max"].flatten()),
+        list(bias_tensor_qsv["min"].flatten()),
+    )
+
+    self.assertIn("StatefulPartitionedCall:0", model_tensor_qsvs)  # output
+    output_qsv = model_tensor_qsvs["StatefulPartitionedCall:0"]
+    self.assertEmpty(output_qsv)
+
+  def test_calibrate_single_fc_success(self):
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    self._calibrator.calibrate(
+        self._representative_dataset, self._recipe_manager
+    )
+    model_tensor_qsvs = self._calibrator.get_model_qsvs()
+
+    self.assertLen(model_tensor_qsvs, 4)
+    self.assertIn("serving_default_input_1:0", model_tensor_qsvs)  # input
+    input_qsv = model_tensor_qsvs["serving_default_input_1:0"]
+    self.assertSequenceAlmostEqual(
+        input_qsv["min"].flatten(), [TEST_MIN_VAL], delta=1e-5
+    )
+    self.assertSequenceAlmostEqual(
+        input_qsv["max"].flatten(), [TEST_MAX_VAL], delta=1e-5
+    )
+
+    self.assertIn("sequential/dense/MatMul", model_tensor_qsvs)  # weight
+    weight_qsv = model_tensor_qsvs["sequential/dense/MatMul"]
+    self.assertSequenceAlmostEqual(weight_qsv["min"].flatten(), [-0.49114203])
+    self.assertSequenceAlmostEqual(weight_qsv["max"].flatten(), [0.4903704])
+
+    self.assertIn(
+        "sequential/dense/BiasAdd/ReadVariableOp", model_tensor_qsvs
+    )  # bias
+    bias_qsv = model_tensor_qsvs["sequential/dense/BiasAdd/ReadVariableOp"]
+    self.assertSequenceAlmostEqual(bias_qsv["min"].flatten(), [-0.38401994])
+    self.assertSequenceAlmostEqual(bias_qsv["max"].flatten(), [0.31727126])
+
+    self.assertIn("StatefulPartitionedCall:0", model_tensor_qsvs)  # output
+    output_qsv = model_tensor_qsvs["StatefulPartitionedCall:0"]
+    # Relu, only check the min
+    self.assertSequenceAlmostEqual(output_qsv["min"].flatten(), [0])
+
+  def test_calibration_cache_is_empty_when_off(self):
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    self.assertEmpty(self._calibrator.get_cached_output())
+    self._calibrator.calibrate(
+        self._representative_dataset, self._recipe_manager, cache_output=False
+    )
+    self.assertEmpty(self._calibrator.get_cached_output())
+
+  def test_calibration_cache_when_on(self):
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    self.assertEmpty(self._calibrator.get_cached_output())
+    self._calibrator.calibrate(
+        self._representative_dataset, self._recipe_manager, cache_output=True
+    )
+    self.assertLen(self._calibrator.get_cached_output(), 10)
+
+  def test_calibration_cache_is_empty_after_reset(self):
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    self._calibrator.calibrate(
+        self._representative_dataset, self._recipe_manager, cache_output=True
+    )
+    self._calibrator.clear_cached_output()
+    self.assertEmpty(self._calibrator.get_cached_output())
+
+  def test_calibrate_unsupported_ops_success(self):
+    # Many ops in the following model are not supported currently.
+    test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/branching_conv_fc.tflite"
+    )
+    test_calibrator = calibrator.Calibrator(test_model_path)
+    _add_default_int8xint8_integer_recipe(self._recipe_manager)
+    dataset_gen = _representative_dataset_gen(size=(3, 4, 4, 1))
+    test_calibrator.calibrate(
+        _get_calibration_data(dataset_gen),
+        self._recipe_manager,
+        cache_output=True,
+    )
+    self.assertLen(test_calibrator.get_cached_output(), 10)
+
+
+class CalibratorAlreadyQuantizedModelTest(googletest.TestCase):
+
+  def test_check_is_float_model_succeeds_when_model_is_float(self):
+    test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/conv_fc_mnist.tflite"
+    )
+    _ = calibrator.Calibrator(test_model_path)
+
+  def test_check_is_float_model_raises_error_when_model_is_quantized(self):
+    test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH, "tests/models/mnist_quantized.tflite"
+    )
+    with self.assertRaisesRegex(
+        ValueError,
+        "The input model for calibration is not a float model.",
+    ):
+      _ = calibrator.Calibrator(test_model_path)
+
+
+class CalibratorToyGemma2Test(googletest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    np.random.seed(0)
+
+    self._test_model_path = os.path.join(
+        TEST_DATA_PREFIX_PATH,
+        "tests/models/toy_model_with_kv_cache_multi_signature.tflite",
+    )
+
+    self._toy_gemma2_calibration_dataset = {
+        "signature_1": [{
+            "cache_0": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+            "cache_1": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+            "positions": _RNG.integers(low=0, high=10, size=(1, 100)).astype(
+                np.int32
+            ),
+            "tokens": _RNG.integers(low=0, high=10, size=(1, 100)).astype(
+                np.int32
+            ),
+        }],
+        "signature_2": [{
+            "cache_0": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+            "cache_1": _RNG.random(size=(1, 100, 4, 4), dtype=np.float32),
+            "positions": _RNG.integers(low=0, high=10, size=(1, 100)).astype(
+                np.int32
+            ),
+            "tokens": _RNG.integers(low=0, high=10, size=(1, 100)).astype(
+                np.int32
+            ),
+        }],
+    }
+
+  def test_toy_gemma2_calibration_success(self):
+    calib = calibrator.Calibrator(self._test_model_path)
+    recipe_mngr = recipe_manager.RecipeManager()
+    _add_default_int8xint8_integer_recipe(recipe_mngr)
+    calib.calibrate(
+        self._toy_gemma2_calibration_dataset,
+        model_recipe_manager=recipe_mngr,
+    )
+    self.assertLen(calib.get_model_qsvs(), 282)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_quantizer/conftest.py

@@ -0,0 +1,22 @@
+# Copyright 2024 The AI Edge Quantizer Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Configuration file for pytest."""
+
+from absl import flags
+
+
+def pytest_configure():
+  flags.FLAGS.mark_as_parsed()