#489 simplified API

nicolay-r · Jul 3, 2023 · 2f9c150 · 2f9c150
1 parent b38995e
commit 2f9c150
Show file tree

Hide file tree

Showing 7 changed files with 15 additions and 29 deletions.
diff --git a/arekit/common/folding/base.py b/arekit/common/folding/base.py
@@ -5,10 +5,6 @@ class BaseDataFolding(object):
         and how to such state into string.
     """
 
-    def __init__(self, supported_data_types=None):
-        assert(isinstance(supported_data_types, list) or supported_data_types is None)
-        self._supported_data_types = supported_data_types
-
     def fold_doc_ids_set(self, doc_ids):
         """ Perform the doc_ids folding process onto provided data_types
         """

diff --git a/arekit/common/folding/nofold.py b/arekit/common/folding/nofold.py
@@ -1,5 +1,3 @@
-import collections
-
 from arekit.common.experiment.data_type import DataType
 from arekit.common.folding.base import BaseDataFolding
 
@@ -8,13 +6,7 @@ class NoFolding(BaseDataFolding):
     """ The case of absent folding in experiment.
     """
 
-    def __init__(self, data_type):
-        assert(isinstance(data_type, DataType))
-        super(NoFolding, self).__init__(supported_data_types=[data_type])
-        self.__data_type = data_type
-
     def fold_doc_ids_set(self, doc_ids):
-        assert(isinstance(doc_ids, collections.Iterable))
-        return {
-            self.__data_type: list(set(doc_ids))
-        }
+        assert(isinstance(doc_ids, dict) and len(doc_ids) == 1)
+        assert(isinstance(list(doc_ids.keys())[0], DataType))
+        return doc_ids
diff --git a/arekit/contrib/utils/cv/two_class.py b/arekit/contrib/utils/cv/two_class.py
@@ -8,13 +8,15 @@ class TwoClassCVFolding(BaseDataFolding):
     """
 
     def __init__(self, supported_data_types, cv_count, splitter):
+        assert(isinstance(supported_data_types, list))
         assert(isinstance(splitter, CrossValidationSplitter))
         assert(isinstance(cv_count, int) and cv_count > 0)
 
         if len(supported_data_types) > 2:
             raise NotImplementedError("Experiments with such amount of data-types are not supported!")
 
-        super(TwoClassCVFolding, self).__init__(supported_data_types=supported_data_types)
+        super(TwoClassCVFolding, self).__init__()
+        self._supported_data_types = supported_data_types
 
         self.__cv_count = cv_count
         self.__splitter = splitter
@@ -59,7 +61,7 @@ def fold_doc_ids_set(self, doc_ids):
             }
 
         if self.__splitter is None:
-            raise NotImplementedError("Splitter has not been intialized!")
+            raise NotImplementedError("Splitter has not been initialized!")
 
         it = self.__splitter.items_to_cv_pairs(doc_ids=set(doc_ids),
                                                cv_count=self.__cv_count)

diff --git a/tests/contrib/utils/test_csv_stream_write.py b/tests/contrib/utils/test_csv_stream_write.py
@@ -65,7 +65,7 @@ def __launch(self, writer, target_extention):
         #####
         # Declaring pipeline related context parameters.
         #####
-        no_folding = NoFolding(data_type=DataType.Train)
+        no_folding = NoFolding()
         doc_provider = FooDocumentProvider()
         text_parser = BaseTextParser(pipeline=[BratTextEntitiesParser(), DefaultTextTokenizer(keep_tokens=True)])
         train_pipeline = text_opinion_extraction_pipeline(
@@ -86,7 +86,7 @@ def __launch(self, writer, target_extention):
                      params_dict={
                          "data_folding": no_folding,
                          "data_type_pipelines": {DataType.Train: train_pipeline},
-                         "doc_ids": [0, 1]
+                         "doc_ids": {DataType.Train: [0, 1]}
                      })
 
     def test_csv_native(self):

diff --git a/tests/tutorials/test_tutorial_data_foldings.py b/tests/tutorials/test_tutorial_data_foldings.py
@@ -26,13 +26,11 @@ def test(self):
             DataType.Test: [4, 5, 6, 7]
         }
 
-        fixed_folding = FixedFolding()
         print("Fixed folding:")
-        self.show_folding(fixed_folding, doc_ids=parts)
+        self.show_folding(FixedFolding(), doc_ids=parts)
 
-        no_folding = NoFolding(data_type=DataType.Train)
         print("No folding:")
-        self.show_folding(no_folding, doc_ids=parts[DataType.Train])
+        self.show_folding(NoFolding(), doc_ids={DataType.Train: parts[DataType.Train]})
 
         splitter_simple = SimpleCrossValidationSplitter(shuffle=True, seed=1)
 

diff --git a/tests/tutorials/test_tutorial_pipeline_sampling_bert.py b/tests/tutorials/test_tutorial_pipeline_sampling_bert.py
@@ -104,7 +104,6 @@ def test(self):
         #####
         # Declaring pipeline related context parameters.
         #####
-        no_folding = NoFolding(data_type=DataType.Train)
         doc_provider = FooDocumentProvider()
         text_parser = BaseTextParser(pipeline=[BratTextEntitiesParser(), DefaultTextTokenizer(keep_tokens=True)])
         train_pipeline = text_opinion_extraction_pipeline(
@@ -123,9 +122,9 @@ def test(self):
 
         pipeline.run(input_data=None,
                      params_dict={
-                         "data_folding": no_folding,
+                         "data_folding": NoFolding(),
                          "data_type_pipelines": {DataType.Train: train_pipeline},
-                         "doc_ids": [0, 1]
+                         "doc_ids": {DataType.Train: [0, 1]}
                      })
 
         reader = PandasCsvReader()

diff --git a/tests/tutorials/test_tutorial_pipeline_sampling_network.py b/tests/tutorials/test_tutorial_pipeline_sampling_network.py
@@ -105,7 +105,6 @@ def test(self):
         #####
         # Declaring pipeline related context parameters.
         #####
-        no_folding = NoFolding(data_type=DataType.Train)
         doc_provider = FooDocumentProvider()
         text_parser = BaseTextParser(pipeline=[
             BratTextEntitiesParser(),
@@ -127,9 +126,9 @@ def test(self):
 
         pipeline.run(input_data=None,
                      params_dict={
-                         "data_folding": no_folding,
+                         "data_folding": NoFolding(),
                          "data_type_pipelines": {DataType.Train: train_pipeline},
-                         "doc_ids": [0, 1]
+                         "doc_ids": {DataType.Train: [0, 1]}
                      })
 
         reader = PandasCsvReader()