add _init_parameters for cuda op

Author: Cathy0908

data_juicer/config/config.py

@@ -425,6 +425,11 @@ def init_setup_from_cfg(cfg: Namespace):
 
     # check number of processes np
     sys_cpu_count = os.cpu_count()
+    if not cfg.np:
+        cfg.np = sys_cpu_count
+        logger.warning(
+            f'Number of processes `np` is not set, '
+            f'Set it to cpu count [{sys_cpu_count}] as default value.')
     if cfg.np > sys_cpu_count:
         logger.warning(f'Number of processes `np` is set as [{cfg.np}], which '
                        f'is larger than the cpu count [{sys_cpu_count}]. Due '

data_juicer/core/ray_data.py

@@ -119,7 +119,14 @@ def _run_single_op(self, op):
                                  1) if op.is_batched_op() else 1
             if isinstance(op, Mapper):
                 if op.use_cuda():
-                    init_params = op._init_parameters
+                    try:
+                        init_params = op._init_parameters
+                    except AttributeError:
+                        raise ValueError(
+                            f'This Op[{op._name}] enables CUDA, you should add'
+                            ' `_init_parameters` attribute to the Op class by '
+                            'add `self._init_parameters = self.remove_extra_parameters(locals())`'  # noqa: E501
+                            ' after super().__init__().')
                     op_args = init_params.pop('args', ())
                     op_kwargs = init_params.pop('kwargs', {})
                     op_kwargs.update(init_params)
@@ -141,7 +148,14 @@ def _run_single_op(self, op):
                                                       num_gpus=num_gpus)
             elif isinstance(op, Filter):
                 if op.use_cuda():
-                    init_params = op._init_parameters
+                    try:
+                        init_params = op._init_parameters
+                    except AttributeError:
+                        raise ValueError(
+                            f'This Op[{op._name}] enables CUDA, you should add'
+                            ' `_init_parameters` attribute to the Op class by '
+                            'add `self._init_parameters = self.remove_extra_parameters(locals())`'  # noqa: E501
+                            ' after super().__init__().')
                     op_args = init_params.pop('args', ())
                     op_kwargs = init_params.pop('kwargs', {})
                     op_kwargs.update(init_params)

data_juicer/ops/filter/image_aesthetics_filter.py

@@ -48,6 +48,7 @@ def __init__(self,
         """
 
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         if hf_scorer_model == '':
             hf_scorer_model = \
                 'shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE'

data_juicer/ops/filter/image_pair_similarity_filter.py

@@ -44,6 +44,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_score = min_score
         self.max_score = max_score
         if any_or_all not in ['any', 'all']:

data_juicer/ops/filter/image_text_matching_filter.py

@@ -53,6 +53,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_score = min_score
         self.max_score = max_score
         if reduce_mode not in ['avg', 'max', 'min']:

data_juicer/ops/filter/image_text_similarity_filter.py

@@ -54,6 +54,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_score = min_score
         self.max_score = max_score
         if reduce_mode not in ['avg', 'max', 'min']:

data_juicer/ops/filter/image_watermark_filter.py

@@ -46,6 +46,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.prob_threshold = prob_threshold
         if any_or_all not in ['any', 'all']:
             raise ValueError(f'Keep strategy [{any_or_all}] is not supported. '

data_juicer/ops/filter/phrase_grounding_recall_filter.py

@@ -115,6 +115,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_recall = min_recall
         self.max_recall = max_recall
         if reduce_mode not in ['avg', 'max', 'min']:

data_juicer/ops/filter/video_aesthetics_filter.py

@@ -75,6 +75,7 @@ def __init__(self,
         """
 
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         if hf_scorer_model == '':
             hf_scorer_model = \
                 'shunk031/aesthetics-predictor-v2-sac-logos-ava1-l14-linearMSE'

data_juicer/ops/filter/video_frames_text_similarity_filter.py

@@ -75,6 +75,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_score = min_score
         self.max_score = max_score
         if frame_sampling_method not in ['all_keyframes', 'uniform']:

data_juicer/ops/filter/video_motion_score_filter.py

@@ -82,6 +82,8 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
+
         self.min_score = min_score
         self.max_score = max_score
         self.sampling_fps = sampling_fps

data_juicer/ops/filter/video_nsfw_filter.py

@@ -66,6 +66,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.score_threshold = score_threshold
         if frame_sampling_method not in ['all_keyframes', 'uniform']:
             raise ValueError(

data_juicer/ops/filter/video_ocr_area_ratio_filter.py

@@ -72,6 +72,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.min_area_ratio = min_area_ratio
         self.max_area_ratio = max_area_ratio
         self.frame_sample_num = frame_sample_num

data_juicer/ops/filter/video_tagging_from_frames_filter.py

@@ -62,6 +62,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         if contain not in ['any', 'all']:
             raise ValueError(f'the containing type [{contain}] is not '
                              f'supported. Can only be one of ["any", "all"].')

data_juicer/ops/filter/video_watermark_filter.py

@@ -70,6 +70,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.prob_threshold = prob_threshold
         if frame_sampling_method not in ['all_keyframes', 'uniform']:
             raise ValueError(

data_juicer/ops/mapper/generate_qa_from_examples_mapper.py

@@ -96,6 +96,7 @@ def __init__(self,
         :param kwargs: Extra keyword arguments.
         """
         super().__init__(**kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
 
         if not seed_file:
             raise ValueError(

data_juicer/ops/mapper/generate_qa_from_text_mapper.py

@@ -69,6 +69,7 @@ def __init__(self,
         """
 
         super().__init__(**kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
 
         if output_pattern is None:
             self.output_pattern = r'Human:(.*?)Assistant:(.*?)(?=Human|$)'  # noqa: E501

data_juicer/ops/mapper/image_tagging_mapper.py

@@ -37,6 +37,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.model_key = prepare_model(
             model_type='recognizeAnything',
             pretrained_model_name_or_path='ram_plus_swin_large_14m.pth',

data_juicer/ops/mapper/optimize_qa_mapper.py

@@ -66,7 +66,7 @@ def __init__(self,
         :param kwargs: Extra keyword arguments.
         """
         super().__init__(**kwargs)
-
+        self._init_parameters = self.remove_extra_parameters(locals())
         self.system_prompt = system_prompt or self.DEFAULT_SYSTEM_PROMPT
         self.input_template = input_template or self.DEFAULT_INPUT_TEMPLATE
         self.qa_pair_template = qa_pair_template or \

data_juicer/ops/mapper/video_captioning_from_audio_mapper.py

@@ -33,6 +33,7 @@ def __init__(self, keep_original_sample: bool = True, *args, **kwargs):
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         AUTOINSTALL.check([
             'transformers', 'transformers_stream_generator', 'einops',
             'accelerate', 'tiktoken'

data_juicer/ops/mapper/video_captioning_from_frames_mapper.py

@@ -109,6 +109,7 @@ def __init__(
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
 
         if keep_candidate_mode not in [
                 'random_any', 'similar_one_simhash', 'all'

data_juicer/ops/mapper/video_captioning_from_summarizer_mapper.py

@@ -82,6 +82,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         AUTOINSTALL.check([
             'torch',
             'transformers',

data_juicer/ops/mapper/video_captioning_from_video_mapper.py

@@ -109,6 +109,7 @@ def __init__(
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
 
         if keep_candidate_mode not in [
                 'random_any', 'similar_one_simhash', 'all'

data_juicer/ops/mapper/video_tagging_from_audio_mapper.py

@@ -38,6 +38,7 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
         AUTOINSTALL.check(['torchaudio'])
         self.model_key = prepare_model(model_type='huggingface',
                                        pretrained_model_name_or_path=hf_ast,

data_juicer/ops/mapper/video_tagging_from_frames_mapper.py

@@ -56,6 +56,8 @@ def __init__(self,
         :param kwargs: extra args
         """
         super().__init__(*args, **kwargs)
+        self._init_parameters = self.remove_extra_parameters(locals())
+
         if frame_sampling_method not in ['all_keyframes', 'uniform']:
             raise ValueError(
                 f'Frame sampling method [{frame_sampling_method}] is not '