From 6b7f79092fcda6b7474949f9873ebb6f19e8cc26 Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <mathieu@feedly.com>
Date: Mon, 22 Mar 2021 19:21:13 -0400
Subject: [PATCH] [DatasetBuilder] union method
---
src/research/common/datasets/dataset_builder.py | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/src/research/common/datasets/dataset_builder.py b/src/research/common/datasets/dataset_builder.py
index 68b595a..d96cebe 100644
--- a/src/research/common/datasets/dataset_builder.py
+++ b/src/research/common/datasets/dataset_builder.py
@@ -9,6 +9,7 @@ from research.common.datasets.lazy_dataset import ExampleT, LazyDataset, TargetT
from research.common.datasets.shuffle_dataset import ShuffleDataset
from research.common.datasets.slice_dataset import SliceDataset
from research.common.datasets.transform_dataset import TransformDataset
+from research.common.datasets.union_dataset import UnionLazyDataset
class DatasetBuilder(Generic[ExampleT, TargetT], Iterable[Tuple[ExampleT, TargetT, str]]):
@@ -71,6 +72,10 @@ class DatasetBuilder(Generic[ExampleT, TargetT], Iterable[Tuple[ExampleT, Target
self.dataset = SliceDataset(self.dataset, start=n)
return self
+ def __or__(self, other: "DatasetBuilder[ExampleT, TargetT]"):
+ self.dataset = UnionLazyDataset((self.dataset, other.dataset))
+ return self
+
@property
def name(self) -> str:
return self.dataset.name
--
GitLab