pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

refine manul parser (#131)

Browse Source
This commit is contained in:
KevinHuSh
2024-03-19 12:26:04 +08:00
committed by GitHub
parent d56c9e7630
commit 9da671b951
13 changed files with 145 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
rag/app/manual.py
View File

@ -2,7 +2,7 @@ import copy
import re
from api.db import ParserType
from rag.nlp import huqie, tokenize, tokenize_table, add_positions
from rag.nlp import huqie, tokenize, tokenize_table, add_positions, bullets_category, title_frequency
from deepdoc.parser import PdfParser
from rag.utils import num_tokens_from_string
@ -14,6 +14,8 @@ class Pdf(PdfParser):
def __call__(self, filename, binary=None, from_page=0,
to_page=100000, zoomin=3, callback=None):
from timeit import default_timer as timer
start = timer()
callback(msg="OCR is running...")
self.__images__(
filename if not binary else binary,
@ -23,19 +25,38 @@ class Pdf(PdfParser):
callback
)
callback(msg="OCR finished.")
#for bb in self.boxes:
# for b in bb:
# print(b)
print("OCR:", timer()-start)
def get_position(bx):
poss = []
pn = bx["page_number"]
top = bx["top"] - self.page_cum_height[pn - 1]
bott = bx["bottom"] - self.page_cum_height[pn - 1]
poss.append((pn, bx["x0"], bx["x1"], top, min(bott, self.page_images[pn-1].size[1]/zoomin)))
while bott * zoomin > self.page_images[pn - 1].size[1]:
bott -= self.page_images[pn- 1].size[1] / zoomin
top = 0
pn += 1
poss.append((pn, bx["x0"], bx["x1"], top, min(bott, self.page_images[pn - 1].size[1] / zoomin)))
return poss
def tag(pn, left, right, top, bottom):
return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##" \
.format(pn, left, right, top, bottom)
from timeit import default_timer as timer
start = timer()
self._layouts_rec(zoomin)
callback(0.65, "Layout analysis finished.")
print("paddle layouts:", timer() - start)
self._table_transformer_job(zoomin)
callback(0.67, "Table analysis finished.")
self._text_merge()
self._concat_downward(concat_between_pages=False)
tbls = self._extract_table_figure(True, zoomin, True, True)
self._naive_vertical_merge()
self._filter_forpages()
callback(0.68, "Text merging finished")
tbls = self._extract_table_figure(True, zoomin, True, True)
# clean mess
for b in self.boxes:
@ -44,25 +65,33 @@ class Pdf(PdfParser):
# merge chunks with the same bullets
self._merge_with_same_bullet()
# merge title with decent chunk
i = 0
while i + 1 < len(self.boxes):
b = self.boxes[i]
if b.get("layoutno","").find("title") < 0:
i += 1
continue
b_ = self.boxes[i + 1]
b_["text"] = b["text"] + "\n" + b_["text"]
b_["x0"] = min(b["x0"], b_["x0"])
b_["x1"] = max(b["x1"], b_["x1"])
b_["top"] = b["top"]
self.boxes.pop(i)
# set pivot using the most frequent type of title,
# then merge between 2 pivot
bull = bullets_category([b["text"] for b in self.boxes])
most_level, levels = title_frequency(bull, [(b["text"], b.get("layout_no","")) for b in self.boxes])
assert len(self.boxes) == len(levels)
sec_ids = []
sid = 0
for i, lvl in enumerate(levels):
if lvl <= most_level: sid += 1
sec_ids.append(sid)
#print(lvl, self.boxes[i]["text"], most_level)
callback(0.8, "Parsing finished")
for b in self.boxes: print(b["text"], b.get("layoutno"))
sections = [(b["text"], sec_ids[i], get_position(b)) for i, b in enumerate(self.boxes)]
for (img, rows), poss in tbls:
sections.append((rows[0], -1, [(p[0]+1, p[1], p[2], p[3], p[4]) for p in poss]))
print(tbls)
return [b["text"] + self._line_tag(b, zoomin) for b in self.boxes], tbls
chunks = []
last_sid = -2
for txt, sec_id, poss in sorted(sections, key=lambda x: (x[-1][0][0], x[-1][0][3], x[-1][0][1])):
poss = "\t".join([tag(*pos) for pos in poss])
if sec_id == last_sid or sec_id == -1:
if chunks:
chunks[-1] += "\n" + txt + poss
continue
chunks.append(txt + poss)
if sec_id >-1: last_sid = sec_id
return chunks
def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", callback=None, **kwargs):
@ -73,7 +102,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
if re.search(r"\.pdf$", filename, re.IGNORECASE):
pdf_parser = Pdf()
cks, tbls = pdf_parser(filename if not binary else binary,
cks = pdf_parser(filename if not binary else binary,
from_page=from_page, to_page=to_page, callback=callback)
else: raise NotImplementedError("file type not supported yet(pdf supported)")
doc = {
@ -84,16 +113,15 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
# is it English
eng = lang.lower() == "english"#pdf_parser.is_english
res = tokenize_table(tbls, doc, eng)
i = 0
chunk = []
tk_cnt = 0
res = []
def add_chunk():
nonlocal chunk, res, doc, pdf_parser, tk_cnt
d = copy.deepcopy(doc)
ck = "\n".join(chunk)
tokenize(d, pdf_parser.remove_tag(ck), pdf_parser.is_english)
tokenize(d, pdf_parser.remove_tag(ck), eng)
d["image"], poss = pdf_parser.crop(ck, need_position=True)
add_positions(d, poss)
res.append(d)
@ -101,7 +129,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
tk_cnt = 0
while i < len(cks):
if tk_cnt > 128: add_chunk()
if tk_cnt > 256: add_chunk()
txt = cks[i]
txt_ = pdf_parser.remove_tag(txt)
i += 1
@ -109,6 +137,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
chunk.append(txt)
tk_cnt += cnt
if chunk: add_chunk()
for i, d in enumerate(res):
print(d)
# d["image"].save(f"./logs/{i}.jpg")
@ -117,6 +146,6 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
if __name__ == "__main__":
import sys
def dummy(a, b):
def dummy(prog=None, msg=""):
pass
chunk(sys.argv[1], callback=dummy)