Source code for data_juicer.ops.mapper.remove_specific_chars_mapper

from typing import List, Union

import regex as re

from ..base_op import OPERATORS, Mapper


@OPERATORS.register_module('remove_specific_chars_mapper')
class RemoveSpecificCharsMapper(Mapper):
    """Mapper to clean specific chars in text samples."""

    _batched_op = True

    def __init__(self,
                 chars_to_remove: Union[str, List[str]] = '◆●■►▼▲▴∆▻▷❖♡□',
                 *args,
                 **kwargs):
        """
        Initialization method.

        :param chars_to_remove: a list or a string including all
            characters that need to be removed from text.
        :param args: extra args
        :param kwargs: extra args
        """

        super().__init__(*args, **kwargs)
        if chars_to_remove:
            self.pattern = '[' + '|'.join(chars_to_remove) + ']'
        else:
            self.pattern = None

    def process_batched(self, samples):
        if self.pattern is None:
            return samples

        samples[self.text_key] = [
            re.sub(pattern=self.pattern,
                   repl=r'',
                   string=text,
                   flags=re.DOTALL) for text in samples[self.text_key]
        ]
        return samples