File size: 5,555 Bytes
f2a3229 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | import logging
import re
import warnings
from bs4 import (
BeautifulSoup,
PageElement,
Tag,
Comment,
Stylesheet,
Script,
NavigableString,
XMLParsedAsHTMLWarning
)
from transformers import BatchEncoding, ProcessorMixin
logger = logging.getLogger(__name__)
warnings.filterwarnings("ignore", category=XMLParsedAsHTMLWarning)
class HTMLLMProcessor(ProcessorMixin):
attributes = ["tokenizer"]
tokenizer_class = "AutoTokenizer"
def __init__(self, tokenizer=None, merge_threshold: int = 2, *args, **kwargs):
super().__init__(tokenizer, *args, **kwargs)
self.tokenizer = tokenizer
self.merge_threshold = merge_threshold
def _process_tree(self, root: Tag) -> Tag:
PRESERVE_TAGS = [
"a", "address", "audio", "br", "button", "canvas", "col", "embed", "figure", "footer",
"form", "frame", "header", "hr", "iframe", "img", "input", "label", "menu",
"meter", "nav", "option", "output", "picture", "progress", "search", "select",
"td", "th", "tr", "textarea", "video", "i", "svg"
]
def _can_tag_be_removed(tag: Tag) -> bool:
is_whitelisted = tag.name in PRESERVE_TAGS
contains_text = tag.get_text(strip=True) != ""
has_children = any(isinstance(c, Tag) for c in tag.children)
return not (is_whitelisted or contains_text or has_children)
try:
# process all children of the tag
i = len(root.contents) - 1
while i >= 0:
child = root.contents[i]
# A: comments, stylesheets, scripts
if isinstance(child, (Comment, Stylesheet, Script)):
# remove the element
child.extract()
# B: tags
elif isinstance(child, Tag):
# first, recursively process the subtree
if len(child.contents) > 0:
self._process_tree(child)
# then, process the tag
if _can_tag_be_removed(child):
# remove only if all the conditions are met
child.extract()
# remove all attributes
child.attrs = {}
# C: texts
elif isinstance(child, NavigableString) and child.get_text(strip=True) == "":
# remove empty text or whitespaces
child.extract()
# D: decrement the index
i -= 1
except RecursionError:
# recursion limit reached -> remove subtree beginning with root
root.extract()
finally:
# remove all attributes
root.attrs = {}
return root
def _merge_tree(self, root: Tag) -> Tag:
def _should_be_merged(tag: Tag) -> bool:
num_children = len(tag.contents)
has_multiple_less_than_threshold = num_children > 1 and num_children < self.merge_threshold
has_single_tag_only = num_children == 1 and (not isinstance(tag.contents[0], NavigableString))
return has_multiple_less_than_threshold or has_single_tag_only
try:
i = 0
# process all childern of the root
while i < len(root.contents):
current_child = root.contents[i]
if isinstance(current_child, Tag): # if child is a HTML tag
if root.name != "html" and _should_be_merged(current_child):
# merge current child children to the root
# remove the current child from the tree
current_child = current_child.extract()
current_child_children = list(current_child.contents)
current_child.clear()
# insert the extracted current child children to the root
for idx, child in enumerate(current_child_children):
root.insert(i + idx, child)
else:
# recursively merge its children
self._merge_tree(current_child)
i += 1
else:
i += 1
except RecursionError:
# recursion limit reached -> remove subtree beginning with root
root.extract()
finally:
return root
def preprocess_html(self, html: str) -> str | None:
try:
root = BeautifulSoup(html, "lxml").html
if not isinstance(root, PageElement):
logger.warning("HTML could not be parsed!")
return None
root = self._process_tree(root)
root = self._merge_tree(root)
html = re.sub(r"\s+", " ", str(root))
html = re.sub(r">\s*<", "><", str(html))
return html
except Exception as e:
logger.warning(f"HTML preprocesing failed: {e}")
return None
def __call__(self, text: str | list[str], *args, **kwargs) -> BatchEncoding:
if isinstance(text, str):
text = self.preprocess_html(text)
elif isinstance(text, list):
text = [self.preprocess_html(t) for t in text]
else:
raise ValueError("HTMLLMProcessor expects the input to be either str or list[str].")
return self.tokenizer(text, *args, **kwargs)
|