fix ray data

Cathy0908 · Sep 9, 2024 · c02a44b · c02a44b
1 parent 1bdb484
commit c02a44b
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/data_juicer/core/ray_data.py b/data_juicer/core/ray_data.py
@@ -301,15 +301,15 @@ def _run_single_op(self, op, dataset: Dataset) -> Dataset:
             batch_size = getattr(op, 'batch_size',
                                  1) if op.is_batched_op() else 1
             if isinstance(op, Mapper):
-                self.data = self.data.map_batches(op.process,
-                                                  batch_size=batch_size,
-                                                  batch_format='pyarrow',
-                                                  num_gpus=num_gpus)
+                dataset = dataset.map_batches(op.process,
+                                              batch_size=batch_size,
+                                              batch_format='pyarrow',
+                                              num_gpus=num_gpus)
             elif isinstance(op, Filter):
-                self.data = self.data.map_batches(op.compute_stats,
-                                                  batch_size=batch_size,
-                                                  batch_format='pyarrow',
-                                                  num_gpus=num_gpus)
+                dataset = dataset.map_batches(op.compute_stats,
+                                              batch_size=batch_size,
+                                              batch_format='pyarrow',
+                                              num_gpus=num_gpus)
                 if op.stats_export_path is not None:
                     dataset.write_json(op.stats_export_path, force_ascii=False)
                 dataset = dataset.filter(op.process)