Sample level error catching

data_juicer/ops/base_op.py

@@ -2,13 +2,76 @@
 
 import pandas as pd
 import pyarrow as pa
+from loguru import logger
 
 from data_juicer.utils.mm_utils import size_to_bytes
 from data_juicer.utils.registry import Registry
 
 OPERATORS = Registry('Operators')
 
 
+def convert_list_dict_to_dict_list(samples):
+    # reconstruct samples from "list of dicts" to "dict of lists"
+    keys = samples[0].keys()
+    res_samples = {}
+    for key in keys:
+        res_samples[key] = [s[key] for s in samples]
+    return res_samples
+
+
+def convert_dict_list_to_list_dict(samples):
+    # reconstruct samples from "dict of lists" to "list of dicts"
+    reconstructed_samples = []
+    keys = list(samples.keys())
+    # take any key, since they should be of same length
+    for i in range(len(samples[keys[0]])):
+        reconstructed_samples.append({key: samples[key][i] for key in samples})
+    return reconstructed_samples
+
+
+def catch_exception_mapper_process(method):
+    """
+    For mapper sample level fault torelerance.
+    """
+
+    def wrapper(self, *args, **kwargs):
+        try:
+            return method(self, *args, **kwargs)
+        except Exception as e:
+            samples = args[0]
+            logger.error(
+                f'An error occurred in mapper operation when processing'
+                f'sample {samples}, {type(e)}: {e}')
+            return {}
+
+    return wrapper
+
+
+def catch_exception_mapper_process_single(method):
+    """
+    For mapper process_single,
+    turn it into batch_size = 1, and enable fault torelerance.
+    """
+
+    def wrapper(self, *args, **kwargs):
+        try:
+            args = list(args)
+            samples = args[0]
+            sample = convert_dict_list_to_list_dict(samples)[0]
+            args[0] = sample
+            args = tuple(args)
+            res_sample = method(self, *args, **kwargs)
+            return convert_list_dict_to_dict_list([res_sample])
+        except Exception as e:
+            samples = args[0]
+            logger.error(
+                f'An error occurred in mapper operation when processing'
+                f'sample {samples}, {type(e)}: {e}')
+            return {}
+
+    return wrapper
+
+
 class OP:
 
     def __init__(self, *args, **kwargs):
@@ -87,6 +150,7 @@ def ray_batch_mapper_wrapper(samples, fn):
     return pa.Table.from_pandas(res)
 
 
+# @mapper_fault_tolerance
 class Mapper(OP):
 
     def __init__(self, *args, **kwargs):
@@ -105,7 +169,9 @@ def __init__(self, *args, **kwargs):
         super(Mapper, self).__init__(*args, **kwargs)
 
         # In default, it's a normal OP instead of batched OP
-        self._batched_op = kwargs.get('batched_op', False)
+        # self._batched_op = kwargs.get('batched_op', False)
+        # Aftet the refactor, we want all ops to be batched OP by default
+        self._batched_op = kwargs.get('batched_op', True)
 
     def process(self, sample):
         """
@@ -134,6 +200,7 @@ def __call__(self, sample):
             return self.process(sample)
 
 
+# @filter_fault_tolerance
 class Filter(OP):
 
     def __init__(self, *args, **kwargs):

data_juicer/ops/mapper/audio_ffmpeg_wrapped_mapper.py

@@ -4,7 +4,7 @@
 from data_juicer.utils.file_utils import transfer_filename
 from data_juicer.utils.logger_utils import HiddenPrints
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 OP_NAME = 'audio_ffmpeg_wrapped_mapper'
 
@@ -47,8 +47,8 @@ def __init__(
         self.capture_stderr = capture_stderr
         self.overwrite_output = overwrite_output
 
-    def process(self, sample):
-        # there is no audio in this sample
+    @catch_exception_mapper_process_single
+    def process(self, sample):  # there is no audio in this sample
         if self.audio_key not in sample or not sample[self.audio_key]:
             return sample
 

data_juicer/ops/mapper/clean_copyright_mapper.py

@@ -4,7 +4,7 @@
 
 import regex as re
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('clean_copyright_mapper')
@@ -23,8 +23,8 @@ def __init__(self, *args, **kwargs):
         self.pat = re.compile('/\\*[^*]*\\*+(?:[^/*][^*]*\\*+)*/')
         self.cpat = re.compile('copyright', re.IGNORECASE)
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
-
         r = self.pat.search(sample[self.text_key])
         if r:
             # found one, now see if it contains "copyright", if so strip it

data_juicer/ops/mapper/clean_email_mapper.py

@@ -1,6 +1,6 @@
 import regex as re
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('clean_email_mapper')
@@ -28,8 +28,8 @@ def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs):
 
         self.repl = repl
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
-
         if not re.search(self.pattern, sample[self.text_key], flags=re.DOTALL):
             return sample
 

data_juicer/ops/mapper/clean_html_mapper.py

@@ -4,7 +4,7 @@
 
 from data_juicer.utils.availability_utils import AvailabilityChecking
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 OP_NAME = 'clean_html_mapper'
 
@@ -25,6 +25,7 @@ def __init__(self, *args, **kwargs):
         """
         super().__init__(*args, **kwargs)
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
 
         def _clean_html(raw_html):

data_juicer/ops/mapper/clean_ip_mapper.py

@@ -1,6 +1,6 @@
 import regex as re
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('clean_ip_mapper')
@@ -32,8 +32,8 @@ def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs):
                 self.pattern = pattern[2:-1]
         self.repl = repl
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
-
         if not re.search(self.pattern, sample[self.text_key], flags=re.DOTALL):
             return sample
 

data_juicer/ops/mapper/clean_links_mapper.py

@@ -3,7 +3,7 @@
 # --------------------------------------------------------
 import regex as re
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('clean_links_mapper')
@@ -38,8 +38,8 @@ def __init__(self, pattern: str = None, repl: str = '', *args, **kwargs):
                 self.pattern = pattern[2:-1]
         self.repl = repl
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
-
         if not re.search(self.pattern, sample[self.text_key], flags=re.DOTALL):
             return sample
 

data_juicer/ops/mapper/expand_macro_mapper.py

@@ -4,7 +4,7 @@
 
 import regex as re
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('expand_macro_mapper')
@@ -55,6 +55,7 @@ def _build_non_arg_macros_dict(self, file_content):
                 macros[macro_name] = macro_val
         return macros
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
         non_arg_macros = self._build_non_arg_macros_dict(sample[self.text_key])
 

data_juicer/ops/mapper/fix_unicode_mapper.py

@@ -1,6 +1,6 @@
 from data_juicer.utils.availability_utils import AvailabilityChecking
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 OP_NAME = 'fix_unicode_mapper'
 
@@ -33,6 +33,7 @@ def __init__(self, normalization: str = None, *args, **kwargs):
                              'supported. Can only be one of '
                              '["NFC", "NFKC", "NFD", "NFKD"]')
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
         sample[self.text_key] = ftfy.fix_text(sample[self.text_key],
                                               normalization=self.normalization)

data_juicer/ops/mapper/image_blur_mapper.py

@@ -6,7 +6,7 @@
 from data_juicer.utils.file_utils import transfer_filename
 from data_juicer.utils.mm_utils import load_data_with_context, load_image
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 from ..op_fusion import LOADED_IMAGES
 
 OP_NAME = 'image_blur_mapper'
@@ -53,6 +53,7 @@ def __init__(self,
         else:
             self.blur = ImageFilter.GaussianBlur(radius)
 
+    @catch_exception_mapper_process_single
     def process(self, sample, context=False):
         # there is no image in this sample
         if self.image_key not in sample or not sample[self.image_key]:

data_juicer/ops/mapper/image_captioning_from_gpt4v_mapper.py

@@ -10,7 +10,7 @@
                                         remove_non_special_tokens,
                                         remove_special_tokens)
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process
 from ..op_fusion import LOADED_IMAGES
 
 SYSTEM_PROMPTS = {
@@ -244,6 +244,7 @@ def _process_single_sample(self, sample):
 
         return [generated_sample]
 
+    @catch_exception_mapper_process
     def process(self, samples):
         # reconstruct samples from "dict of lists" to "list of dicts"
         reconstructed_samples = []

data_juicer/ops/mapper/image_captioning_mapper.py

@@ -13,7 +13,7 @@
                                         remove_special_tokens)
 from data_juicer.utils.model_utils import get_model, prepare_model
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process
 from ..op_fusion import LOADED_IMAGES
 
 OP_NAME = 'image_captioning_mapper'
@@ -272,6 +272,7 @@ def _reduce_captions_per_image(self, chunk,
                 generated_text_candidates_single_chunk[max_index])
         return new_generated_text_per_chunk
 
+    @catch_exception_mapper_process
     def process(self, samples, rank=None):
         """
         Note:

data_juicer/ops/mapper/image_diffusion_mapper.py

@@ -10,7 +10,7 @@
                                         load_image, remove_special_tokens)
 from data_juicer.utils.model_utils import get_model, prepare_model
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process
 from ..op_fusion import LOADED_IMAGES
 
 OP_NAME = 'image_diffusion_mapper'
@@ -209,6 +209,7 @@ def _process_single_sample(self, ori_sample, rank=None, context=False):
 
         return generated_samples
 
+    @catch_exception_mapper_process
     def process(self, samples, rank=None, context=False):
         """
             Note:

data_juicer/ops/mapper/image_face_blur_mapper.py

@@ -6,7 +6,7 @@
 from data_juicer.utils.mm_utils import (load_data_with_context, load_image,
                                         pil_to_opencv)
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 from ..op_fusion import LOADED_IMAGES
 
 OP_NAME = 'image_face_blur_mapper'
@@ -66,6 +66,7 @@ def __init__(self,
         # Initialize face detector
         self.detector = dlib.get_frontal_face_detector()
 
+    @catch_exception_mapper_process_single
     def process(self, sample, context=False):
         # there is no image in this sample
         if self.image_key not in sample or not sample[self.image_key]:

data_juicer/ops/mapper/nlpaug_en_mapper.py

@@ -4,7 +4,7 @@
 
 from data_juicer.utils.availability_utils import AvailabilityChecking
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process
 
 OP_NAME = 'nlpaug_en_mapper'
 
@@ -122,6 +122,7 @@ def __init__(self,
         else:
             self.aug = aug_pipeline
 
+    @catch_exception_mapper_process
     def process(self, samples):
         # no augmentation methods are opened
         if len(self.aug) == 0:

data_juicer/ops/mapper/nlpcda_zh_mapper.py

@@ -5,7 +5,7 @@
 from data_juicer.utils.availability_utils import AvailabilityChecking
 from data_juicer.utils.logger_utils import HiddenPrints
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process
 
 OP_NAME = 'nlpcda_zh_mapper'
 
@@ -128,6 +128,7 @@ def __init__(self,
                 self.aug_pipeline.append(
                     nlpcda.EquivalentChar(create_num=create_num))
 
+    @catch_exception_mapper_process
     def process(self, samples):
         # no augmentation methods are opened
         if len(self.aug_pipeline) == 0:

data_juicer/ops/mapper/punctuation_normalization_mapper.py

@@ -2,7 +2,7 @@
 # https://github.com/bigscience-workshop/data-preparation
 # --------------------------------------------------------
 
-from ..base_op import OPERATORS, Mapper
+from ..base_op import OPERATORS, Mapper, catch_exception_mapper_process_single
 
 
 @OPERATORS.register_module('punctuation_normalization_mapper')
@@ -55,6 +55,7 @@ def __init__(self, *args, **kwargs):
             '►': '-',
         }
 
+    @catch_exception_mapper_process_single
     def process(self, sample):
         sample[self.text_key] = ''.join([
             self.punctuation_unicode.get(c, c) for c in sample[self.text_key]