Introduce DatasetWithSeqId

Author: nikita-savelyevv

src/nncf/__init__.py

@@ -18,6 +18,7 @@
 from nncf.common.strip import strip as strip
 from nncf.config import NNCFConfig as NNCFConfig
 from nncf.data import Dataset as Dataset
+from nncf.data import DatasetWithSeqId as DatasetWithSeqId
 from nncf.errors import BufferFullError as BufferFullError
 from nncf.errors import InstallationError as InstallationError
 from nncf.errors import InternalError as InternalError

src/nncf/data/__init__.py

@@ -10,4 +10,5 @@
 # limitations under the License.
 
 from nncf.data.dataset import Dataset as Dataset
+from nncf.data.dataset import DatasetWithSeqId as DatasetWithSeqId
 from nncf.data.generators import generate_text_data as generate_text_data

src/nncf/data/dataset.py

@@ -38,12 +38,6 @@ class Dataset:
         will be passed into the model as-is.
     """
 
-    # RESET_STATE_KEY is a special input key used by OpenVINO backend to control resetting of internal model state
-    # between model inferences. This key can be added to a dataset sample input dictionary with either
-    # `True` or `False` value. With `True` value, the model state will be reset before inference on the corresponding
-    # sample, and with `False` the state will not be reset.
-    RESET_STATE_KEY = "reset_state"
-
     def __init__(self, data_source: Iterable[Any], transform_func: Optional[Callable[..., Any]] = None):
         self._data_source = data_source
         self._transform_func = transform_func
@@ -94,6 +88,41 @@ def get_batch_size(self) -> Optional[int]:
         return None
 
 
+@api(canonical_alias="nncf.DatasetWithSeqId")
+class DatasetWithSeqId(Dataset):
+    def __init__(
+        self,
+        data_source: Iterable[Any],
+        sequence_ids: Iterable[int],
+        transform_func: Optional[Callable[..., Any]] = None,
+    ):
+        """
+        A dataset wrapper that associates each data item with a sequence ID. Sequence IDs are used to reset state
+        of stateful models between different sequences during inference.
+        :param data_source: The iterable object serving as the source of data items.
+        :param sequence_ids: The iterable of sequence IDs corresponding to each data item in the data source. Must
+            have the same length as the data source and contain only integers.
+        :param transform_func: The function that is used to extract the model's input
+            from the data item. The data item here is the data item that is returned from
+            the data source per iteration. This function should be passed when
+            the data item cannot be directly used as model's input. If this is not specified, then the data item
+            will be passed into the model as-is.
+        """
+        if any(not isinstance(it, int) for it in sequence_ids):
+            msg = "All sequence IDs must be integers."
+            raise ValueError(msg)
+        if transform_func is not None:
+            transform_func = lambda x, seq_id: (transform_func(x), seq_id)
+        super().__init__(data_source, transform_func)
+        self._sequence_ids = sequence_ids
+
+    def get_data(self, indices: Optional[list[int]] = None) -> Iterable[Any]:
+        return DataProvider(zip(self._data_source, self._sequence_ids), None, indices)
+
+    def get_inference_data(self, indices: Optional[list[int]] = None) -> Iterable[Any]:
+        return DataProvider(zip(self._data_source, self._sequence_ids), self._transform_func, indices)
+
+
 class DataProvider:
     def __init__(
         self,

src/nncf/openvino/engine.py

@@ -16,7 +16,6 @@
 from openvino import Type
 from openvino.properties.hint import inference_precision
 
-import nncf
 from nncf.common.engine import Engine
 from nncf.openvino.graph.model_utils import model_has_state
 
@@ -35,7 +34,11 @@ def __init__(self, compiled_model: ov.CompiledModel, stateful: bool):
         self.reset_state = stateful and hasattr(self.infer_request, "reset_state")
 
     def infer(
-        self, input_data: Union[np.ndarray, list[np.ndarray], tuple[np.ndarray], dict[str, np.ndarray]]
+        self,
+        input_data: Union[
+            Union[np.ndarray, list[np.ndarray], tuple[np.ndarray], dict[str, np.ndarray]],
+            tuple[Union[np.ndarray, list[np.ndarray], tuple[np.ndarray], dict[str, np.ndarray]], int],
+        ],
     ) -> dict[str, np.ndarray]:
         """
         Runs model on the provided input via OpenVINO Runtime.
@@ -44,10 +47,11 @@ def infer(
         :param input_data: Inputs for the model.
         :return output_data: Model's output.
         """
-        if isinstance(input_data, dict) and nncf.Dataset.RESET_STATE_KEY in input_data:
-            # In this case state resetting is controlled by the input data flag
-            input_data = input_data.copy()
-            if input_data.pop(nncf.Dataset.RESET_STATE_KEY):
+        if isinstance(input_data, tuple) and len(input_data) == 2 and isinstance(input_data[1], int):
+            # A dataset with sequence ids is provided
+            input_data, seq_id = input_data
+            if seq_id == 0:
+                # Reset state only at the beginning of a new sequence
                 if self.reset_state:
                     self.infer_request.reset_state()
                 else:

src/nncf/version.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "3.0.0"
+__version__ = "3.0.0.dev0+37b9d3650dirty"
 
 
 BKC_TORCH_SPEC = "==2.9.*"

tests/openvino/native/test_engine.py

@@ -95,31 +95,30 @@ def _reset_state():
 
     model = StatefulModel(True).ov_model
     inp = model.get_parameters()[0]
-    input_data = [{"input_data": np.ones(inp.shape), nncf.Dataset.RESET_STATE_KEY: False} for _ in range(10)]
-    reset_ind = [2, 5, 7]
-    for ind in reset_ind:
-        input_data[ind][nncf.Dataset.RESET_STATE_KEY] = True
+    input_data = [{"input_data": np.ones(inp.shape)} for _ in range(10)]
+    # sequence lengths are [2, 4, 4]
+    dataset = nncf.DatasetWithSeqId(input_data, [0, 1, 0, 1, 2, 3, 0, 1, 2, 3])
 
     engine = OVNativeEngine(model)
     reset_order = []
     wrap_reset_state(engine.engine.infer_request)
 
-    for inp_data in input_data:
+    for inp_data in dataset.get_inference_data():
         engine.infer(inp_data)
         reset_order.append("infer")
 
     assert reset_order == [
+        "reset",
         "infer",
         "infer",
         "reset",
         "infer",
         "infer",
         "infer",
-        "reset",
-        "infer",
         "infer",
         "reset",
         "infer",
         "infer",
         "infer",
+        "infer",
     ]