Add stop_on_shortest flag

facebookresearch · Oct 2, 2023 · a24baa0 · a24baa0
1 parent 73e348f
commit a24baa0
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 6 deletions.
diff --git a/fairseq2n/src/fairseq2n/data/data_pipeline.cc b/fairseq2n/src/fairseq2n/data/data_pipeline.cc
@@ -228,7 +228,8 @@ data_pipeline::round_robin(std::vector<data_pipeline> pipelines)
 data_pipeline_builder
 data_pipeline::sample(
     std::vector<data_pipeline> pipelines,
-    std::optional<std::vector<float32>> weights)
+    std::optional<std::vector<float32>> weights,
+    bool stop_at_shortest)
 {
     if (pipelines.empty())
         throw_<std::invalid_argument>(
@@ -251,8 +252,8 @@ data_pipeline::sample(
 
     auto tmp = std::make_shared<std::vector<data_pipeline>>(std::move(pipelines));
 
-    auto factory = [tmp, weights=std::move(weights.value())]() mutable {
-        return std::make_unique<sample_data_source>(std::move(*tmp), std::move(weights));
+    auto factory = [tmp, weights=std::move(weights.value()), stop_at_shortest]() mutable {
+        return std::make_unique<sample_data_source>(std::move(*tmp), std::move(weights), stop_at_shortest);
     };
 
     return data_pipeline_builder{std::move(factory)};

diff --git a/fairseq2n/src/fairseq2n/data/data_pipeline.h b/fairseq2n/src/fairseq2n/data/data_pipeline.h
@@ -86,7 +86,8 @@ class FAIRSEQ2_API data_pipeline {
     static data_pipeline_builder
     sample(
         std::vector<data_pipeline> pipelines,
-        std::optional<std::vector<float>> weights = {});
+        std::optional<std::vector<float>> weights = {},
+        bool stop_at_shortest = true);
 
     static data_pipeline_builder
     constant(data example, std::optional<std::string> key = {});

diff --git a/fairseq2n/src/fairseq2n/data/sample_data_source.cc b/fairseq2n/src/fairseq2n/data/sample_data_source.cc
@@ -16,7 +16,7 @@
 
 namespace fairseq2n::detail {
 
-sample_data_source::sample_data_source(std::vector<data_pipeline> &&pipelines, std::vector<float32> &&weights)
+sample_data_source::sample_data_source(std::vector<data_pipeline> &&pipelines, std::vector<float32> &&weights, bool stop_at_shortest)
     : pipelines_(std::move(pipelines))
 {
     weights_ = make_tensor_from_vector(weights, { static_cast<std::int64_t>(pipelines_.size()) });

diff --git a/fairseq2n/src/fairseq2n/data/sample_data_source.h b/fairseq2n/src/fairseq2n/data/sample_data_source.h
@@ -21,7 +21,7 @@ namespace fairseq2n::detail {
 class sample_data_source final : public data_source {
 public:
     explicit
-    sample_data_source(std::vector<data_pipeline> &&pipelines, std::vector<float32> &&weights);
+    sample_data_source(std::vector<data_pipeline> &&pipelines, std::vector<float32> &&weights, bool stop_at_shortest);
 
     std::optional<data>
     next() override;

diff --git a/src/fairseq2/data/data_pipeline.py b/src/fairseq2/data/data_pipeline.py
@@ -105,13 +105,16 @@ def round_robin(pipelines: Sequence["DataPipeline"]) -> "DataPipelineBuilder":
         def sample(
             pipelines: Sequence["DataPipeline"],
             weights: Optional[Sequence[float]] = None,
+            stop_at_shortest: bool = True,
         ) -> "DataPipelineBuilder":
             """Extract examples from ``pipelines`` by sampling based on ``weights``.
 
             :param data_pipelines:
                 The data pipelines to sample from.
             :param weights:
                 Desired distribution of pipelines. If None, use uniform distribution.
+            :param stop_at_shortest:
+                Flag to stop sampling when first pipeline reaches its end.
             """
 
         @staticmethod