#63 now we consider to keep labels for any split

nicolay-r · Aug 8, 2023 · 3623170 · 3623170
1 parent fe0fe25
commit 3623170
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/arekit_ss/framework/arekit/serialize_bert.py b/arekit_ss/framework/arekit/serialize_bert.py
@@ -1,6 +1,5 @@
 from arekit.common.data.const import ENTITIES, ENTITY_TYPES
 from arekit.common.data.input.providers.rows.samples import BaseSampleRowProvider
-from arekit.common.experiment.data_type import DataType
 from arekit.contrib.utils.data.storages.row_cache import RowCacheStorage
 
 from arekit.contrib.utils.io_utils.samples import SamplesIO
@@ -15,5 +14,5 @@ def serialize_bert_pipeline(writer, rows_provider, output_dir):
         storage=RowCacheStorage(force_collect_columns=[ENTITIES, ENTITY_TYPES]),
         balance_func=lambda _: False,
         samples_io=SamplesIO(target_dir=output_dir, writer=writer),
-        save_labels_func=lambda data_type: data_type != DataType.Test,
+        save_labels_func=lambda _: True,
         rows_provider=rows_provider)
diff --git a/arekit_ss/framework/arekit/serialize_nn.py b/arekit_ss/framework/arekit/serialize_nn.py
@@ -1,5 +1,4 @@
 from arekit.common.data.const import ENTITIES, ENTITY_TYPES
-from arekit.common.experiment.data_type import DataType
 from arekit.contrib.networks.input.const import FrameVariantIndices, FrameConnotations
 from arekit.contrib.utils.data.storages.row_cache import RowCacheStorage
 from arekit.contrib.utils.data.writers.base import BaseWriter
@@ -23,5 +22,5 @@ def serialize_nn_pipeline(output_dir, writer, rows_provider):
         samples_io=SamplesIO(target_dir=output_dir, writer=writer),
         emb_io=NpEmbeddingIO(target_dir=output_dir),
         balance_func=lambda _: False,
-        save_labels_func=lambda data_type: data_type != DataType.Test,
-        save_embedding=True)
+        save_labels_func=lambda _: True,
+        save_embedding=True)