Source code for data_juicer.ops.mapper.expand_macro_mapper

# Some code here has been modified from:
# https://github.com/togethercomputer/RedPajama-Data/blob/rp_v1/data_prep/arxiv/arxiv_cleaner.py
# --------------------------------------------------------

import regex as re

from ..base_op import OPERATORS, Mapper




@OPERATORS.register_module('expand_macro_mapper')
class ExpandMacroMapper(Mapper):
    """Mapper to expand macro definitions in the document body of Latex
    samples."""

    _batched_op = True



    def __init__(self, *args, **kwargs):
        """
        Initialization method.

        :param args: extra args
        :param kwargs: extra args
        """
        super().__init__(*args, **kwargs)


    def _build_non_arg_macros_dict(self, file_content):
        # regex for extracting \newcommand macros without arguments
        non_arg_nc_reg = re.compile(
            # this regex matches the following:
            # \newcommand{\macro_name}{macro_value}
            # \newcommand*{\macro_name}{macro_value}
            # where macro_name is only allowed to contain letters and numbers;
            # macro_value can contain any character.
            pattern=r'\\\bnewcommand\b\*?\{(\\[a-zA-Z0-9]+?)\}\{(.*?)\}$',
            flags=re.MULTILINE)

        # regex for extracting \def macros without arguments
        non_arg_def_reg = re.compile(
            # this regex matches the following:
            # \def\macro_name{macro_value}
            # where macro_name is only allowed to contain letters and numbers;
            # macro_value can contain any character.
            pattern=r'\\def\s*(\\[a-zA-Z0-9]+?)\s*\{(.*?)\}$',
            flags=re.MULTILINE)

        # Extract all user-defined LaTeX macros from the preamble
        macros = {}
        for reg in [non_arg_nc_reg, non_arg_def_reg]:
            for match in reg.finditer(file_content):
                # convert the macro name and value to a raw string that can be
                # used in re.sub
                macro_name = match.group(1).encode('unicode-escape').decode(
                    'utf-8')
                macro_val = match.group(2).encode('unicode-escape').decode(
                    'utf-8')

                macros[macro_name] = macro_val
        return macros



    def process_batched(self, samples):
        for idx, text in enumerate(samples[self.text_key]):
            non_arg_macros = self._build_non_arg_macros_dict(text)

            # TODO: macros that take arguments are not supported yet
            arg_macros = {}

            # inline-expand all non-arg macros
            for macro_name, macro_value in non_arg_macros.items():
                text = re.sub(
                    # make pattern grouped to make sure that the macro
                    # is not part of a longer alphanumeric word
                    pattern=r'(' + macro_name + r')' + r'([^a-zA-Z0-9])',
                    # replace the macro with its value and add back the
                    # character that was matched after the macro
                    repl=macro_value + r'\2',
                    string=text)

            # inline-expand all macros that use args
            # TODO: inline-expand macros with args
            for macro_name, macro_value in arg_macros.items():
                pass

            samples[self.text_key][idx] = text

        return samples