Implement bucketed weighted sampling for VITS (#1871)

tests/data_tests/test_samplers.py

@@ -5,11 +5,11 @@ import unittest
 import torch
 
 from TTS.config.shared_configs import BaseDatasetConfig
-from TTS.encoder.utils.samplers import PerfectBatchSampler
 from TTS.tts.datasets import load_tts_samples
 from TTS.tts.utils.data import get_length_balancer_weights
 from TTS.tts.utils.languages import get_language_balancer_weights
 from TTS.tts.utils.speakers import get_speaker_balancer_weights
+from TTS.utils.samplers import BucketBatchSampler, PerfectBatchSampler
 
 # Fixing random state to avoid random fails
 torch.manual_seed(0)
@@ -163,3 +163,31 @@ class TestSamplers(unittest.TestCase):
                 else:
                     len2 += 1
             assert is_balanced(len1, len2), "Length Weighted sampler is supposed to be balanced"
+
+    def test_bucket_batch_sampler(self):
+        bucket_size_multiplier = 2
+        sampler = range(len(train_samples))
+        sampler = BucketBatchSampler(
+            sampler,
+            data=train_samples,
+            batch_size=7,
+            drop_last=True,
+            sort_key=lambda x: len(x["text"]),
+            bucket_size_multiplier=bucket_size_multiplier,
+        )
+
+        # check if the samples are sorted by text lenght whuile bucketing
+        min_text_len_in_bucket = 0
+        bucket_items = []
+        for batch_idx, batch in enumerate(list(sampler)):
+            if (batch_idx + 1) % bucket_size_multiplier == 0:
+                for bucket_item in bucket_items:
+                    self.assertLessEqual(min_text_len_in_bucket, len(train_samples[bucket_item]["text"]))
+                    min_text_len_in_bucket = len(train_samples[bucket_item]["text"])
+                min_text_len_in_bucket = 0
+                bucket_items = []
+            else:
+                bucket_items += batch
+
+        # check sampler length
+        self.assertEqual(len(sampler), len(train_samples) // 7)