Repository: lhotse-speech/lhotse

Issue:
BucketingSampler with equal_duration drops an arbitrary(?) number of cuts.
If you use the BucketingSampler with num_buckets, especially with even number of buckets, though I think you can also get it to work for an odd number of buckets, at the very least the last cut is dropped, but if you choose the durations and num_buckets appropriate you can have many more dropped.

The reason this is not caught in the test is that the line

`assert set(cut_set.ids) == set(c.id for c in sampled_cuts)`

is only called for the equal_length method and not the equal_duration method.

The problem arises from the fact that the loops in _create_buckets_equal_duration_single() are not over the cuts, but are instead over the buckets and the duration over the buckets. Because you pop the last cut off of the odd buckets and put it back at the front, you can be left with a non-empty iterator. 

At the very least the BucketingSampler equal_duration feature should be updated with a warning when using it and print out the number of cuts it has dropped. Many times it will be 0, or 1. 

However, if modify

`def test_bucketing_sampler_single_cuts():
    cut_set = DummyManifest(CutSet, begin_id=0, end_id=1000)
    sampler = BucketingSampler(cut_set, sampler_type=SimpleCutSampler, bucket_method="equal_duration", num_buckets=500)
    print(f'Num cuts orig: {len(cut_set)}, Num cuts in sampler: {sum(len(b[0]) for b in sampler.buckets)}')
    sampled_cuts = []
    for batch in sampler:
        sampled_cuts.extend(batch)
    assert set(cut_set.ids) == set(c.id for c in sampled_cuts)`

in test/dataset/sampling/test_sampling.py to look like what you have above, you will see that in that example it drops 25% of the cuts.



Comments:

Patch:
diff --git a/lhotse/dataset/sampling/bucketing.py b/lhotse/dataset/sampling/bucketing.py
index 04541b37..a2ad7ba3 100644
--- a/lhotse/dataset/sampling/bucketing.py
+++ b/lhotse/dataset/sampling/bucketing.py
@@ -422,23 +422,55 @@ def _create_buckets_equal_duration_single(
     """
     total_duration = np.sum(c.duration for c in cuts)
     bucket_duration = total_duration / num_buckets
-    iter_cuts = iter(cuts)
-    buckets = []
-    for bucket_idx in range(num_buckets):
-        bucket = []
-        current_duration = 0
-        try:
-            while current_duration < bucket_duration:
-                bucket.append(next(iter_cuts))
-                current_duration += bucket[-1].duration
-            # Every odd bucket, take the cut that exceeded the bucket's duration
-            # and put it in the front of the iterable, so that it goes to the
-            # next bucket instead. It will ensure that the last bucket is not too
-            # thin (otherwise all the previous buckets are a little too large).
-            if bucket_idx % 2:
-                last_cut = bucket.pop()
-                iter_cuts = chain([last_cut], iter_cuts)
-        except StopIteration:
-            assert bucket_idx == num_buckets - 1
-        buckets.append(CutSet.from_cuts(bucket))
+    # Define the order for adding cuts. We start at the beginning, then go to
+    # the end, and work our way to the middle. Once in the middle we distribute
+    # excess cuts among the two buckets close to the median duration. This
+    # handles the problem of where to place cuts that caused previous buckets
+    # to "over-flow" without sticking all of them in the last bucket, which
+    # causes one large bucket at the end and also places many small duration
+    # cuts with longer ones.
+    order = list(range(0, len(cuts), 2)) + list(
+        range(len(cuts) - (1 + len(cuts) % 2), 0, -2)
+    )
+    order2idx = {o_idx: i for i, o_idx in enumerate(order)}
+    durations = [c.duration for c in cuts]
+
+    # We need a list of the cut durations in the same order (0, N-1, 1, N-2, ...)
+    ordered_cut_durations = sorted(zip(order, durations), key=lambda x: x[0])
+    last_order, first_bucket = 0, 0
+    last_bucket = num_buckets - 1
+    buckets_dict = {i: 0 for i in range(num_buckets)}
+    buckets_cut_dict = {i: [] for i in range(num_buckets)}
+    middle_bucket = None
+    idx_to_bucket_id = {}
+    for i, (order_idx, duration) in enumerate(ordered_cut_durations, 1):
+        # Check if we are at the middle bucket. first_bucket is the left bucket
+        # we are processing. last_bucket is the right bucket. When they are the
+        # same we are filling the bucket with cuts near the median duration.
+        if middle_bucket is None and first_bucket == last_bucket:
+            middle_bucket = first_bucket
+
+        # i % 2 = 1 ==> process the left_bucket (first_bucket)
+        if i % 2:
+            if buckets_dict[first_bucket] + duration > bucket_duration:
+                if middle_bucket is not None and first_bucket == middle_bucket:
+                    first_bucket = min(middle_bucket - 1, num_buckets - 1)
+                else:
+                    first_bucket = min(first_bucket + 1, num_buckets - 1)
+            buckets_dict[first_bucket] += duration
+            idx_to_bucket_id[order2idx[order_idx]] = first_bucket
+        # i % 2 = 0 ==> process the right bucket (last_bucket)
+        else:
+            if buckets_dict[last_bucket] + duration > bucket_duration:
+                if middle_bucket is not None and last_bucket == middle_bucket:
+                    last_bucket = max(middle_bucket + 1, 0)
+                else:
+                    last_bucket = max(last_bucket - 1, 0)
+            buckets_dict[last_bucket] += duration
+            idx_to_bucket_id[order2idx[order_idx]] = last_bucket
+
+    # Now that buckets have been assigned, create the new cutset.
+    for cut_idx, cut in enumerate(cuts):
+        buckets_cut_dict[idx_to_bucket_id[cut_idx]].append(cut)
+    buckets = [CutSet.from_cuts(buckets_cut_dict[i]) for i in range(num_buckets)]
     return buckets


Test patch:
diff --git a/test/dataset/sampling/test_sampling.py b/test/dataset/sampling/test_sampling.py
index 4c922d40..1ed453ef 100644
--- a/test/dataset/sampling/test_sampling.py
+++ b/test/dataset/sampling/test_sampling.py
@@ -467,19 +467,29 @@ def test_bucketing_sampler_single_cuts_equal_duration():
     )
 
     # Ensure that each consecutive bucket has less cuts than the previous one
-    prev_len = float("inf")
-    bucket_cum_durs = []
+    sampled_cuts, bucket_cum_durs = [], []
+    prev_min, prev_max = 0, 0
+    num_overlapping_bins = 0
     for (bucket,) in sampler.buckets:
-        bucket_cum_durs.append(sum(c.duration for c in bucket))
-        curr_len = len(bucket)
-        assert curr_len < prev_len
-        prev_len = curr_len
+        bucket_durs = [c.duration for c in bucket]
+        sampled_cuts.extend(c for c in bucket)
+        bucket_cum_durs.append(sum(bucket_durs))
+        bucket_min, bucket_max = min(bucket_durs), max(bucket_durs)
+        # Ensure that bucket lengths do not overlap, except for the middle
+        # 3 buckets maybe
+        if prev_max > bucket_min:
+            num_overlapping_bins += 1
+        assert num_overlapping_bins < 3
+        prev_min = bucket_min
+        prev_max = bucket_max
 
     # Assert that all bucket cumulative durations are within 1/10th of the mean
     mean_bucket_dur = mean(bucket_cum_durs)  # ~ 1300s
     for d in bucket_cum_durs:
         assert abs(d - mean_bucket_dur) < 0.1 * mean_bucket_dur
 
+    assert set(cut_set.ids) == set(c.id for c in sampled_cuts)
+
 
 def test_bucketing_sampler_shuffle():
     cut_set = DummyManifest(CutSet, begin_id=0, end_id=10)