* modified according to the latest buffer imp.

Author: HYLcool

tests/data/controllers/task_parser_test.py

@@ -43,11 +43,6 @@ def _run_test(self, rft_config, return_none=False):
             logger.info("None dj config.")
         else:
             self.assertIsNotNone(dj_config)
-            op_weights = {}
-            for op in dj_config.process:
-                op_name = list(op.keys())[0]
-                op_weights[op_name] = op[op_name]["op_weight"]
-            logger.info(op_weights)
 
     def test_instruction1(self):
         rft_config = DataPipelineConfig()

tests/data/core/dataset_test.py

@@ -57,6 +57,7 @@ def test_rft_dataset_init(self):
         dataset = RftDataset(
             data_pipeline_config=self.data_pipeline_config, reward_schema="default"
         )
+        dataset.read_from_buffer()
 
         self.assertEqual(len(dataset), 10)
         self.assertIsInstance(dataset.reward_schema, RewardSchema)
@@ -65,6 +66,7 @@ def test_format_dataset(self):
         dataset = RftDataset(
             data_pipeline_config=self.data_pipeline_config, reward_schema="default"
         )
+        dataset.read_from_buffer()
         original_data = dataset.data
         # no formatter
         dataset.format(formatters=[])

tests/data/core/formatter_test.py

@@ -54,6 +54,7 @@ def test_init(self):
 
     def test_transform(self):
         dataset = RftDataset(data_pipeline_config=self.data_config, reward_schema="default")
+        dataset.read_from_buffer()
         formatter = BoxedMathAnswerFormatter(config=self.data_config.format)
         self.assertNotIn(formatter.config.response_key, dataset.data.column_names)
         dataset.format(formatter)

tests/data/processor/cleaner_test.py

@@ -41,6 +41,7 @@ def _run_test(self, tgt_list, weight=1, data_dist="gaussian"):
         )
 
         dataset = RftDataset(self.rft_config.data_processor.task_pipeline)
+        dataset.read_from_buffer()
         dataset = cleaner.process([dataset])
 
         res_list = dataset.data.select_columns("text").to_list()

trinity/buffer/reader/file_reader.py

@@ -235,10 +235,14 @@ def read(self, strategy: Optional[ReadStrategy] = None):
 @FILE_READERS.register_module("raw")
 class RawDataReader(BufferReader):
     def __init__(self, meta: StorageConfig, config: BufferConfig):
+        self.returned = False
         self.dataset = load_dataset(meta.path, name=meta.subset_name, split=meta.split)
 
     def __len__(self):
         return len(self.dataset)
 
     def read(self, strategy: Optional[ReadStrategy] = None) -> List:
+        if self.returned:
+            raise StopIteration
+        self.returned = True
         return self.dataset.to_list()

trinity/buffer/writer/file_writer.py

@@ -22,7 +22,7 @@ def __init__(self, meta: StorageConfig, config: BufferConfig):
         ext = os.path.splitext(meta.path)[-1]
         if ext != ".jsonl":
             raise ValueError(f"File path must end with .json or .jsonl, got {meta.path}")
-        self.writer = jl.open(meta.path, mode="w")
+        self.writer = jl.open(meta.path, mode="a")
 
     def write(self, data: List) -> None:
         self.writer.write_all(data)

trinity/data/controllers/active_iterator.py

@@ -118,43 +118,53 @@ def run(self):
             traceback.print_exc()
             return 3, "DataCleaner loading failed."
 
-        # step 4. apply processors to calculate scores of different dimensions
-        try:
-            res_dataset = dataset
-            if hit_cleaner:
-                res_dataset = cleaner.process([res_dataset])
-            if hit_synthesizer:
-                res_dataset = synthesizer.process([res_dataset])
-            if hit_human_annotator:
-                res_dataset = human_annotator.process([res_dataset])
-        except Exception:
-            traceback.print_exc()
-            return 4, "DataProcessors processing failed."
-
-        # step 5. calculate the average and final scores, including priority
-        try:
-            if hit_cleaner:
-                scored_dataset = self._group_scores(res_dataset)
-                scored_dataset = self._compute_priority_scores(scored_dataset)
-            else:
-                scored_dataset = res_dataset
-        except Exception:
-            traceback.print_exc()
-            return 5, "Grouping and computing priority score failed."
-
-        # step 6. track lineage if they are changed
-        try:
-            res_dataset = scored_dataset
-        except Exception:
-            traceback.print_exc()
-            return 6, "Tracking lineage failed."
-
-        # step 7. export the result to the output buffer
-        try:
-            res_dataset.write_to_buffer()
-        except Exception:
-            traceback.print_exc()
-            return 7, "Exporting result to output buffer failed."
+        while True:
+            # step 4. load data from the input buffers for the next batch
+            try:
+                dataset.read_from_buffer()
+            except StopIteration:
+                break
+            except Exception:
+                traceback.print_exc()
+                return 4, "RftDataset loading from buffers failed."
+
+            # step 5. apply processors to calculate scores of different dimensions
+            try:
+                res_dataset = dataset
+                if hit_cleaner:
+                    res_dataset = cleaner.process([res_dataset])
+                if hit_synthesizer:
+                    res_dataset = synthesizer.process([res_dataset])
+                if hit_human_annotator:
+                    res_dataset = human_annotator.process([res_dataset])
+            except Exception:
+                traceback.print_exc()
+                return 5, "DataProcessors processing failed."
+
+            # step 6. calculate the average and final scores, including priority
+            try:
+                if hit_cleaner:
+                    scored_dataset = self._group_scores(res_dataset)
+                    scored_dataset = self._compute_priority_scores(scored_dataset)
+                else:
+                    scored_dataset = res_dataset
+            except Exception:
+                traceback.print_exc()
+                return 6, "Grouping and computing priority score failed."
+
+            # step 7. track lineage if they are changed
+            try:
+                res_dataset = scored_dataset
+            except Exception:
+                traceback.print_exc()
+                return 7, "Tracking lineage failed."
+
+            # step 8. export the result to the output buffer
+            try:
+                res_dataset.write_to_buffer()
+            except Exception:
+                traceback.print_exc()
+                return 8, "Exporting result to output buffer failed."
 
         return 0, "success"
 

trinity/data/core/dataset.py

@@ -45,11 +45,10 @@ def __init__(
         input_buffer_configs = self.config.input_buffers
         if len(input_buffer_configs) == 0:
             raise ValueError("input_buffers is empty in data pipeline config")
-        datasets = []
+        self.buffers = []
         for input_buffer_config in input_buffer_configs:
-            input_buffer = get_buffer_reader(input_buffer_config, self.buffer_config)
-            datasets.append(Dataset.from_list(input_buffer.read()))
-        self.data = concatenate_datasets(datasets)
+            self.buffers.append(get_buffer_reader(input_buffer_config, self.buffer_config))
+        self.data = Dataset.from_list([])
 
         self.reward_schema = self._init_reward_schema(reward_schema)
         self.stats: Dict[str, Any] = {}
@@ -65,6 +64,12 @@ def format(
         for formatter in formatters:
             self.data = formatter(self.data, num_proc)
 
+    def read_from_buffer(self):
+        datasets = []
+        for buffer in self.buffers:
+            datasets.append(Dataset.from_list(buffer.read()))
+        self.data = concatenate_datasets(datasets)
+
     def write_to_buffer(
         self, output_storage_config: StorageConfig = None, buffer_config: BufferConfig = None
     ):
@@ -75,6 +80,7 @@ def write_to_buffer(
         output_buffer = get_buffer_writer(output_storage_config, buffer_config)
         output_buffer.write(self.data.to_list())
         output_buffer.finish()
+        self.data = Dataset.from_list([])
 
     def to_parquet(self, path: str):
         self.data.to_parquet(path)