Add non-implemented batch functions

iterative · May 8, 2022 · bfce0c0 · bfce0c0
1 parent fca1b89
commit bfce0c0
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 2 deletions.
diff --git a/mlem/contrib/lightgbm.py b/mlem/contrib/lightgbm.py
@@ -1,7 +1,7 @@
 import os
 import posixpath
 import tempfile
-from typing import Any, ClassVar, List, Optional, Tuple, Type
+from typing import Any, ClassVar, Iterator, List, Optional, Tuple, Type
 
 import lightgbm as lgb
 from pydantic import BaseModel
@@ -114,6 +114,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             )
         )
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class LightGBMModelIO(ModelIO):
     """

diff --git a/mlem/contrib/numpy.py b/mlem/contrib/numpy.py
@@ -188,6 +188,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             data = self.dataset_type.actual_type(res)
             return self.dataset_type.copy().bind(data)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class NumpyArrayWriter(DatasetWriter):
     """DatasetWriter implementation for numpy ndarray"""

diff --git a/mlem/contrib/pandas.py b/mlem/contrib/pandas.py
@@ -598,6 +598,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             data.index.name = None
         return self.dataset_type.copy().bind(data)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class PandasSeriesWriter(DatasetWriter, _PandasIO):
     """DatasetWriter for pandas series"""

diff --git a/mlem/contrib/torch.py b/mlem/contrib/torch.py
@@ -1,4 +1,4 @@
-from typing import Any, ClassVar, Optional, Tuple
+from typing import Any, ClassVar, Iterator, Optional, Tuple
 
 import torch
 
@@ -100,6 +100,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             data = torch.load(f)
             return self.dataset_type.copy().bind(data)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class TorchModelIO(ModelIO):
     """

diff --git a/mlem/core/dataset_type.py b/mlem/core/dataset_type.py
@@ -218,6 +218,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
                 data = self.dataset_type.to_type(res)
             return self.dataset_type.copy().bind(data)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class ListDatasetType(DatasetType, DatasetSerializer):
     """
@@ -295,6 +300,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             data_list.append(elem_dtype.data)
         return self.dataset_type.copy().bind(data_list)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class _TupleLikeDatasetType(DatasetType, DatasetSerializer):
     """
@@ -393,6 +403,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
         data_list = self.dataset_type.actual_type(data_list)
         return self.dataset_type.copy().bind(data_list)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 class TupleLikeListDatasetType(_TupleLikeDatasetType):
     """
@@ -548,6 +563,11 @@ def read(self, artifacts: Artifacts) -> DatasetType:
             data_dict[key] = v_dataset_type.data
         return self.dataset_type.copy().bind(data_dict)
 
+    def read_batch(
+        self, artifacts: Artifacts, batch: int
+    ) -> Iterator[DatasetType]:
+        raise NotImplementedError
+
 
 #
 #