add _accelerator

Cathy0908 · Jul 24, 2024 · 149023c · 149023c
1 parent 3496228
commit 149023c
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 1 deletion.
diff --git a/data_juicer/ops/mapper/extract_qa_mapper.py b/data_juicer/ops/mapper/extract_qa_mapper.py
@@ -25,7 +25,6 @@ class ExtractQAMapper(Mapper):
     """
 
     _accelerator = 'cuda'
-    _batched_op = True
 
     def __init__(self,
                  hf_model: str = 'alibaba-pai/pai-qwen1_5-7b-doc2qa',

diff --git a/data_juicer/ops/mapper/generate_instruction_mapper.py b/data_juicer/ops/mapper/generate_instruction_mapper.py
@@ -38,6 +38,7 @@ class GenerateInstructionMapper(Mapper):
     The number of samples generated is determined by
     the length of the empty dataset.
     """
+    _accelerator = 'cuda'
 
     def __init__(self,
                  hf_model,

diff --git a/data_juicer/ops/mapper/optimize_instruction_mapper.py b/data_juicer/ops/mapper/optimize_instruction_mapper.py
@@ -16,6 +16,7 @@ class OptimizeInstructionMapper(Mapper):
         alibaba-pai/Qwen2-7B-Instruct-Refine
     ]
     """
+    _accelerator = 'cuda'
 
     def __init__(self,
                  hf_model: str = 'alibaba-pai/Qwen2-7B-Instruct-Refine',