+ add model_params to text_tagging_by_prompt_mapper

+ flush the buffer when outputting the trace results and wait for 1 sec

Author: HYLcool

data_juicer/core/tracer/ray_tracer.py

@@ -167,5 +167,7 @@ def finalize_traces(self):
             # We'll use a generic name for now, could be improved with operator type detection
             res_name = self.get_trace_file_path(op_name)
             dif_df = pd.DataFrame(traces)
-            dif_df.to_json(res_name, orient="records", lines=True, force_ascii=False)
+            with open(res_name, "w") as out_buf:
+                dif_df.to_json(out_buf, orient="records", lines=True, force_ascii=False)
+                out_buf.flush()
             print(f"Exported {len(traces)} traced samples for op [{op_name}] to {res_name}")

data_juicer/ops/mapper/text_tagging_by_prompt_mapper.py

@@ -74,6 +74,7 @@ def __init__(
         tensor_parallel_size: int = None,
         max_model_len: int = None,
         max_num_seqs: int = 256,
+        model_params: Dict = None,
         sampling_params: Dict = None,
         *args,
         **kwargs,
@@ -93,6 +94,7 @@ def __init__(
             derived from the model config.
         :param max_num_seqs: It is only valid when enable_vllm is True.
             Maximum number of sequences to be processed in a single iteration.
+        :param model_params: Parameters for model initialization.
         :param sampling_params: Sampling parameters for text generation.
             e.g {'temperature': 0.9, 'top_p': 0.95}
         :param args: extra args
@@ -117,7 +119,9 @@ def __init__(
         self.prompt = prompt
         self.tag_list = tag_list
         self.enable_vllm = enable_vllm
-        model_params = {"trust_remote_code": trust_remote_code, "max_num_seqs": max_num_seqs}
+        if model_params is None:
+            model_params = {}
+        model_params.update({"trust_remote_code": trust_remote_code, "max_num_seqs": max_num_seqs})
         if tensor_parallel_size is not None:
             model_params["tensor_parallel_size"] = tensor_parallel_size
         if max_model_len is not None:

tests/core/tracer/test_ray_tracer.py

@@ -2,6 +2,7 @@
 import unittest
 import tempfile
 import shutil
+import time
 import jsonlines as jl
 from data_juicer.core.tracer.ray_tracer import RayTracer
 from data_juicer.utils.unittest_utils import TEST_TAG
@@ -58,6 +59,7 @@ def test_collect_mapper_sample_basic(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_mapper.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -87,6 +89,7 @@ def test_collect_mapper_sample_no_change(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_mapper.jsonl')
         # File should not exist since no samples were collected
@@ -105,6 +108,7 @@ def test_collect_mapper_sample_with_trace_keys(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_mapper.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -135,6 +139,7 @@ def test_collect_mapper_sample_with_missing_trace_keys(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_mapper.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -166,6 +171,7 @@ def test_collect_mapper_sample_not_in_op_list(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_mapper.jsonl')
         self.assertFalse(os.path.exists(trace_file_path))
@@ -183,6 +189,7 @@ def test_collect_filter_sample_basic(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_filter.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -208,6 +215,7 @@ def test_collect_filter_sample_should_keep(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_filter.jsonl')
         self.assertFalse(os.path.exists(trace_file_path))
@@ -225,6 +233,7 @@ def test_collect_filter_sample_not_in_op_list(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-test_filter.jsonl')
         self.assertFalse(os.path.exists(trace_file_path))
@@ -254,6 +263,7 @@ def test_collect_mapper_sample_show_num_limit(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-limited_mapper.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -288,6 +298,7 @@ def test_collect_filter_sample_show_num_limit(self):
 
         # Finalize traces to write to file
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         trace_file_path = os.path.join(self.work_dir, 'trace', 'sample_trace-limited_filter.jsonl')
         self.assertTrue(os.path.exists(trace_file_path))
@@ -327,6 +338,7 @@ def test_finalize_traces_empty(self):
 
         # Don't collect anything, just finalize
         ray.get(tracer.finalize_traces.remote())
+        time.sleep(1)
 
         # No trace files should exist
         trace_dir = os.path.join(self.work_dir, 'trace')

tests/ops/mapper/test_text_tagging_by_prompt_mapper.py

@@ -52,7 +52,9 @@ def test_tagging_vllm(self):
             enable_vllm=True,
             max_model_len=1024,
             max_num_seqs=16,
-            sampling_params={'temperature': 0.1, 'top_p': 0.95, 'max_tokens': 256})
+            sampling_params={'temperature': 0.1, 'top_p': 0.95, 'max_tokens': 256},
+            model_params={'gpu_memory_utilization': 0.8},
+        )
 
 
 if __name__ == '__main__':