Source code for data_juicer.ops.mapper.remove_header_mapper
# Some code here has been modified from:
# https://github.com/togethercomputer/RedPajama-Data/tree/rp_v1/
# --------------------------------------------------------
import regex as re
from ..base_op import OPERATORS, Mapper
[docs]@OPERATORS.register_module('remove_header_mapper')
class RemoveHeaderMapper(Mapper):
"""Mapper to remove headers at the beginning of documents in Latex
samples."""
_batched_op = True
[docs] def __init__(self, drop_no_head: bool = True, *args, **kwargs):
"""
Initialization method.
:param drop_no_head: whether to drop sample texts without
headers.
:param args: extra args
:param kwargs: extra args
"""
super().__init__(*args, **kwargs)
self.pattern = r'^(.*?)('
self.pattern += r'\\\bchapter\b\*?(?:\[(.*?)\])?\{(.*?)\}|'
self.pattern += r'\\\bpart\b\*?(?:\[(.*?)\])?\{(.*?)\}|'
self.pattern += r'\\\bsection\b\*?(?:\[(.*?)\])?\{(.*?)\}|'
self.pattern += r'\\\bsubsection\b\*?(?:\[(.*?)\])?\{(.*?)\}|'
self.pattern += r'\\\bsubsubsection\b\*?(?:\[(.*?)\])?\{(.*?)\}|'
self.pattern += r'\\\bparagraph\b\*?(?:\[(.*?)\])?\{(.*?)\}'
self.pattern += r'\\\bsubparagraph\b\*?(?:\[(.*?)\])?\{(.*?)\}'
self.pattern += r')'
self.drop_no_head = drop_no_head
[docs] def process_batched(self, samples):
for idx, text in enumerate(samples[self.text_key]):
if not re.search(self.pattern, text, flags=re.DOTALL):
if self.drop_no_head:
text = ''
continue
text = re.sub(pattern=self.pattern,
repl=r'\2',
string=text,
flags=re.DOTALL)
samples[self.text_key][idx] = text
return samples