fix(dataset): skip datasets allotted no samples under max_sample_num (#1095)

* fix(dataset): honor a 0-sample share under max_sample_num

* fix(dataset): skip datasets allotted no samples under max_sample_num

Author: moneebullah25

data_juicer/core/data/dataset_builder.py

@@ -173,6 +173,11 @@ def load_dataset(self, **kwargs) -> DJDataset:
         _datasets = []
         # load datasets with sample numbers
         for stra, weight, sample_num in zip(self.load_strategies, self.weights, self.sample_numbers):
+            # random_sample treats 0 as "unset" and falls back to the raw weight,
+            # so a dataset allotted no samples must be skipped, not sampled
+            if self.max_sample_num and sample_num <= 0:
+                continue
+
             # load dataset with its load strategy
             dataset = stra.load_data(**load_kwargs)
 

tests/core/data/test_dataset_builder.py

@@ -520,6 +520,26 @@ def test_invalid_max_sample_num(self):
             self.assertIn('should be a positive integer', 
                           str(context.exception))
 
+    def test_max_sample_num_zero_share(self):
+        """Test max_sample_num with datasets allotted 0 samples"""
+        cfg = Namespace()
+        cfg.dataset = {
+            'configs': [
+                {
+                    'type': 'local',
+                    'path': 'test_data/sample.jsonl',
+                    'weight': 1.0
+                } for _ in range(4)
+            ],
+            'max_sample_num': 2
+        }
+
+        builder = DatasetBuilder(cfg)
+        self.assertEqual(builder.sample_numbers, [1, 1, 0, 0])
+
+        ds = builder.load_dataset()
+        self.assertEqual(len(ds), 2)
+
     @TEST_TAG('ray')
     def test_builder_ray_config(self):
         """Test loading Ray configuration from YAML"""