mirror of
https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git
synced 2025-04-23 22:50:17 +08:00
add paper & manual parser (#46)
This commit is contained in:
parent
04aba1bb65
commit
96a1a44cb6
@ -1,5 +1,9 @@
|
|||||||
import re
|
import re
|
||||||
|
|
||||||
|
from nltk import word_tokenize
|
||||||
|
|
||||||
|
from rag.nlp import stemmer, huqie
|
||||||
|
|
||||||
|
|
||||||
def callback__(progress, msg, func):
|
def callback__(progress, msg, func):
|
||||||
if not func :return
|
if not func :return
|
||||||
@ -46,3 +50,21 @@ def bullets_category(sections):
|
|||||||
res = i
|
res = i
|
||||||
maxium = h
|
maxium = h
|
||||||
return res
|
return res
|
||||||
|
|
||||||
|
def is_english(texts):
|
||||||
|
eng = 0
|
||||||
|
for t in texts:
|
||||||
|
if re.match(r"[a-zA-Z]", t.strip()):
|
||||||
|
eng += 1
|
||||||
|
if eng / len(texts) > 0.8:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def tokenize(d, t, eng):
|
||||||
|
d["content_with_weight"] = t
|
||||||
|
if eng:
|
||||||
|
t = re.sub(r"([a-z])-([a-z])", r"\1\2", t)
|
||||||
|
d["content_ltks"] = " ".join([stemmer.stem(w) for w in word_tokenize(t)])
|
||||||
|
else:
|
||||||
|
d["content_ltks"] = huqie.qie(t)
|
||||||
|
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
|
@ -3,12 +3,13 @@ import re
|
|||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from docx import Document
|
from docx import Document
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from rag.app import callback__, bullets_category, BULLET_PATTERN
|
from rag.app import callback__, bullets_category, BULLET_PATTERN, is_english, tokenize
|
||||||
from rag.nlp import huqie
|
from rag.nlp import huqie
|
||||||
|
from rag.parser.docx_parser import HuDocxParser
|
||||||
from rag.parser.pdf_parser import HuParser
|
from rag.parser.pdf_parser import HuParser
|
||||||
|
|
||||||
|
|
||||||
class Docx(object):
|
class Docx(HuDocxParser):
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@ -42,14 +43,7 @@ class Pdf(HuParser):
|
|||||||
print("paddle layouts:", timer()-start)
|
print("paddle layouts:", timer()-start)
|
||||||
bxs = self.sort_Y_firstly(self.boxes, np.median(self.mean_height) / 3)
|
bxs = self.sort_Y_firstly(self.boxes, np.median(self.mean_height) / 3)
|
||||||
# is it English
|
# is it English
|
||||||
eng = 0
|
eng = is_english([b["text"] for b in bxs])
|
||||||
for b in bxs:
|
|
||||||
if re.match(r"[a-zA-Z]", b["text"].strip()):
|
|
||||||
eng += 1
|
|
||||||
if eng / len(bxs) > 0.8:
|
|
||||||
eng = True
|
|
||||||
else:
|
|
||||||
eng = False
|
|
||||||
# Merge vertically
|
# Merge vertically
|
||||||
i = 0
|
i = 0
|
||||||
while i + 1 < len(bxs):
|
while i + 1 < len(bxs):
|
||||||
@ -59,7 +53,7 @@ class Pdf(HuParser):
|
|||||||
bxs.pop(i)
|
bxs.pop(i)
|
||||||
continue
|
continue
|
||||||
concatting_feats = [
|
concatting_feats = [
|
||||||
b["text"].strip()[-1] in ",;:'\",、‘“;:",
|
b["text"].strip()[-1] in ",;:'\",、‘“;:-",
|
||||||
len(b["text"].strip())>1 and b["text"].strip()[-2] in ",;:'\",‘“、;:",
|
len(b["text"].strip())>1 and b["text"].strip()[-2] in ",;:'\",‘“、;:",
|
||||||
b["text"].strip()[0] in "。;?!?”)),,、:",
|
b["text"].strip()[0] in "。;?!?”)),,、:",
|
||||||
]
|
]
|
||||||
@ -118,14 +112,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
|
|||||||
sections = [l for l in sections if l]
|
sections = [l for l in sections if l]
|
||||||
|
|
||||||
# is it English
|
# is it English
|
||||||
eng = 0
|
eng = is_english(sections)
|
||||||
for sec in sections:
|
|
||||||
if re.match(r"[a-zA-Z]", sec.strip()):
|
|
||||||
eng += 1
|
|
||||||
if eng / len(sections) > 0.8:
|
|
||||||
eng = True
|
|
||||||
else:
|
|
||||||
eng = False
|
|
||||||
# Remove 'Contents' part
|
# Remove 'Contents' part
|
||||||
i = 0
|
i = 0
|
||||||
while i < len(sections):
|
while i < len(sections):
|
||||||
@ -181,8 +168,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
|
|||||||
if pdf_parser:
|
if pdf_parser:
|
||||||
d["image"] = pdf_parser.crop(ck)
|
d["image"] = pdf_parser.crop(ck)
|
||||||
ck = pdf_parser.remove_tag(ck)
|
ck = pdf_parser.remove_tag(ck)
|
||||||
d["content_ltks"] = huqie.qie(ck)
|
tokenize(d, ck, eng)
|
||||||
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
|
|
||||||
res.append(d)
|
res.append(d)
|
||||||
return res
|
return res
|
||||||
|
|
||||||
|
140
rag/app/manual.py
Normal file
140
rag/app/manual.py
Normal file
@ -0,0 +1,140 @@
|
|||||||
|
import copy
|
||||||
|
import re
|
||||||
|
from collections import Counter
|
||||||
|
from rag.app import callback__, bullets_category, BULLET_PATTERN, is_english, tokenize
|
||||||
|
from rag.nlp import huqie, stemmer
|
||||||
|
from rag.parser.docx_parser import HuDocxParser
|
||||||
|
from rag.parser.pdf_parser import HuParser
|
||||||
|
from nltk.tokenize import word_tokenize
|
||||||
|
import numpy as np
|
||||||
|
from rag.utils import num_tokens_from_string
|
||||||
|
|
||||||
|
|
||||||
|
class Pdf(HuParser):
|
||||||
|
def __call__(self, filename, binary=None, from_page=0,
|
||||||
|
to_page=100000, zoomin=3, callback=None):
|
||||||
|
self.__images__(
|
||||||
|
filename if not binary else binary,
|
||||||
|
zoomin,
|
||||||
|
from_page,
|
||||||
|
to_page)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: OCR finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
|
||||||
|
from timeit import default_timer as timer
|
||||||
|
start = timer()
|
||||||
|
self._layouts_paddle(zoomin)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Layout analysis finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
print("paddle layouts:", timer() - start)
|
||||||
|
self._table_transformer_job(zoomin)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Table analysis finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
self._text_merge()
|
||||||
|
column_width = np.median([b["x1"] - b["x0"] for b in self.boxes])
|
||||||
|
self._concat_downward(concat_between_pages=False)
|
||||||
|
self._filter_forpages()
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Text merging finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
tbls = self._extract_table_figure(True, zoomin, False)
|
||||||
|
|
||||||
|
# clean mess
|
||||||
|
for b in self.boxes:
|
||||||
|
b["text"] = re.sub(r"([\t ]|\u3000){2,}", " ", b["text"].strip())
|
||||||
|
|
||||||
|
# merge chunks with the same bullets
|
||||||
|
i = 0
|
||||||
|
while i + 1 < len(self.boxes):
|
||||||
|
b = self.boxes[i]
|
||||||
|
b_ = self.boxes[i + 1]
|
||||||
|
if b["text"].strip()[0] != b_["text"].strip()[0] \
|
||||||
|
or b["page_number"]!=b_["page_number"] \
|
||||||
|
or b["top"] > b_["bottom"]:
|
||||||
|
i += 1
|
||||||
|
continue
|
||||||
|
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)
|
||||||
|
# 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)
|
||||||
|
|
||||||
|
for b in self.boxes: print(b["text"], b.get("layoutno"))
|
||||||
|
|
||||||
|
print(tbls)
|
||||||
|
return [b["text"] + self._line_tag(b, zoomin) for b in self.boxes], tbls
|
||||||
|
|
||||||
|
|
||||||
|
def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
|
||||||
|
pdf_parser = None
|
||||||
|
paper = {}
|
||||||
|
|
||||||
|
if re.search(r"\.pdf$", filename, re.IGNORECASE):
|
||||||
|
pdf_parser = Pdf()
|
||||||
|
cks, tbls = pdf_parser(filename if not binary else binary,
|
||||||
|
from_page=from_page, to_page=to_page, callback=callback)
|
||||||
|
doc = {
|
||||||
|
"docnm_kwd": filename
|
||||||
|
}
|
||||||
|
doc["title_tks"] = huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", doc["docnm_kwd"]))
|
||||||
|
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
||||||
|
# is it English
|
||||||
|
eng = pdf_parser.is_english
|
||||||
|
|
||||||
|
res = []
|
||||||
|
# add tables
|
||||||
|
for img, rows in tbls:
|
||||||
|
bs = 10
|
||||||
|
de = ";" if eng else ";"
|
||||||
|
for i in range(0, len(rows), bs):
|
||||||
|
d = copy.deepcopy(doc)
|
||||||
|
r = de.join(rows[i:i + bs])
|
||||||
|
r = re.sub(r"\t——(来自| in ).*”%s" % de, "", r)
|
||||||
|
tokenize(d, r, eng)
|
||||||
|
d["image"] = img
|
||||||
|
res.append(d)
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
chunk = []
|
||||||
|
tk_cnt = 0
|
||||||
|
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)
|
||||||
|
d["image"] = pdf_parser.crop(ck)
|
||||||
|
res.append(d)
|
||||||
|
chunk = []
|
||||||
|
tk_cnt = 0
|
||||||
|
|
||||||
|
while i < len(cks):
|
||||||
|
if tk_cnt > 128: add_chunk()
|
||||||
|
txt = cks[i]
|
||||||
|
txt_ = pdf_parser.remove_tag(txt)
|
||||||
|
i += 1
|
||||||
|
cnt = num_tokens_from_string(txt_)
|
||||||
|
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")
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import sys
|
||||||
|
|
||||||
|
chunk(sys.argv[1])
|
240
rag/app/paper.py
Normal file
240
rag/app/paper.py
Normal file
@ -0,0 +1,240 @@
|
|||||||
|
import copy
|
||||||
|
import re
|
||||||
|
from collections import Counter
|
||||||
|
from rag.app import callback__, bullets_category, BULLET_PATTERN, is_english, tokenize
|
||||||
|
from rag.nlp import huqie, stemmer
|
||||||
|
from rag.parser.docx_parser import HuDocxParser
|
||||||
|
from rag.parser.pdf_parser import HuParser
|
||||||
|
from nltk.tokenize import word_tokenize
|
||||||
|
import numpy as np
|
||||||
|
from rag.utils import num_tokens_from_string
|
||||||
|
|
||||||
|
|
||||||
|
class Pdf(HuParser):
|
||||||
|
def __call__(self, filename, binary=None, from_page=0,
|
||||||
|
to_page=100000, zoomin=3, callback=None):
|
||||||
|
self.__images__(
|
||||||
|
filename if not binary else binary,
|
||||||
|
zoomin,
|
||||||
|
from_page,
|
||||||
|
to_page)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: OCR finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
|
||||||
|
from timeit import default_timer as timer
|
||||||
|
start = timer()
|
||||||
|
self._layouts_paddle(zoomin)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Layout analysis finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
print("paddle layouts:", timer() - start)
|
||||||
|
self._table_transformer_job(zoomin)
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Table analysis finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
self._text_merge()
|
||||||
|
column_width = np.median([b["x1"] - b["x0"] for b in self.boxes])
|
||||||
|
self._concat_downward(concat_between_pages=False)
|
||||||
|
self._filter_forpages()
|
||||||
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page / 4,
|
||||||
|
"Page {}~{}: Text merging finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
tbls = self._extract_table_figure(True, zoomin, False)
|
||||||
|
|
||||||
|
# clean mess
|
||||||
|
if column_width < self.page_images[0].size[0] / zoomin / 2:
|
||||||
|
print("two_column...................", column_width,
|
||||||
|
self.page_images[0].size[0] / zoomin / 2)
|
||||||
|
self.boxes = self.sort_X_by_page(self.boxes, column_width / 2)
|
||||||
|
for b in self.boxes:
|
||||||
|
b["text"] = re.sub(r"([\t ]|\u3000){2,}", " ", b["text"].strip())
|
||||||
|
freq = Counter([b["text"] for b in self.boxes])
|
||||||
|
garbage = set([k for k, v in freq.items() if v > self.total_page * 0.6])
|
||||||
|
i = 0
|
||||||
|
while i < len(self.boxes):
|
||||||
|
if self.boxes[i]["text"] in garbage \
|
||||||
|
or (re.match(r"[a-zA-Z0-9]+$", self.boxes[i]["text"]) and not self.boxes[i].get("layoutno")) \
|
||||||
|
or (i + 1 < len(self.boxes) and self.boxes[i]["text"] == self.boxes[i + 1]["text"]):
|
||||||
|
self.boxes.pop(i)
|
||||||
|
elif i + 1 < len(self.boxes) and self.boxes[i].get("layoutno", '0') == self.boxes[i + 1].get("layoutno",
|
||||||
|
'1'):
|
||||||
|
# merge within same layouts
|
||||||
|
self.boxes[i + 1]["top"] = self.boxes[i]["top"]
|
||||||
|
self.boxes[i + 1]["x0"] = min(self.boxes[i]["x0"], self.boxes[i + 1]["x0"])
|
||||||
|
self.boxes[i + 1]["x1"] = max(self.boxes[i]["x1"], self.boxes[i + 1]["x1"])
|
||||||
|
self.boxes[i + 1]["text"] = self.boxes[i]["text"] + " " + self.boxes[i + 1]["text"]
|
||||||
|
self.boxes.pop(i)
|
||||||
|
else:
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
def _begin(txt):
|
||||||
|
return re.match(
|
||||||
|
"[0-9. 一、i]*(introduction|abstract|摘要|引言|keywords|key words|关键词|background|背景|目录|前言|contents)",
|
||||||
|
txt.lower().strip())
|
||||||
|
|
||||||
|
# get title and authors
|
||||||
|
title = ""
|
||||||
|
authors = []
|
||||||
|
i = 0
|
||||||
|
while i < min(32, len(self.boxes)):
|
||||||
|
b = self.boxes[i]
|
||||||
|
i += 1
|
||||||
|
if b.get("layoutno", "").find("title") >= 0:
|
||||||
|
title = b["text"]
|
||||||
|
if _begin(title):
|
||||||
|
title = ""
|
||||||
|
break
|
||||||
|
for j in range(3):
|
||||||
|
if _begin(self.boxes[i + j]["text"]): break
|
||||||
|
authors.append(self.boxes[i + j]["text"])
|
||||||
|
break
|
||||||
|
break
|
||||||
|
# get abstract
|
||||||
|
abstr = ""
|
||||||
|
i = 0
|
||||||
|
while i + 1 < min(32, len(self.boxes)):
|
||||||
|
b = self.boxes[i]
|
||||||
|
i += 1
|
||||||
|
txt = b["text"].lower().strip()
|
||||||
|
if re.match("(abstract|摘要)", txt):
|
||||||
|
if len(txt.split(" ")) > 32 or len(txt) > 64:
|
||||||
|
abstr = txt + self._line_tag(b, zoomin)
|
||||||
|
i += 1
|
||||||
|
break
|
||||||
|
txt = self.boxes[i + 1]["text"].lower().strip()
|
||||||
|
if len(txt.split(" ")) > 32 or len(txt) > 64:
|
||||||
|
abstr = txt + self._line_tag(self.boxes[i + 1], zoomin)
|
||||||
|
i += 1
|
||||||
|
break
|
||||||
|
if not abstr: i = 0
|
||||||
|
|
||||||
|
for b in self.boxes: print(b["text"], b.get("layoutno"))
|
||||||
|
print(tbls)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"title": title if title else filename,
|
||||||
|
"authors": " ".join(authors),
|
||||||
|
"abstract": abstr,
|
||||||
|
"lines": [(b["text"] + self._line_tag(b, zoomin), b.get("layoutno", "")) for b in self.boxes[i:] if
|
||||||
|
re.match(r"(text|title)", b.get("layoutno", "text"))],
|
||||||
|
"tables": tbls
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
|
||||||
|
pdf_parser = None
|
||||||
|
paper = {}
|
||||||
|
|
||||||
|
if re.search(r"\.pdf$", filename, re.IGNORECASE):
|
||||||
|
pdf_parser = Pdf()
|
||||||
|
paper = pdf_parser(filename if not binary else binary,
|
||||||
|
from_page=from_page, to_page=to_page, callback=callback)
|
||||||
|
doc = {
|
||||||
|
"docnm_kwd": paper["title"] if paper["title"] else filename,
|
||||||
|
"authors_tks": paper["authors"]
|
||||||
|
}
|
||||||
|
doc["title_tks"] = huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", doc["docnm_kwd"]))
|
||||||
|
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
||||||
|
doc["authors_sm_tks"] = huqie.qieqie(doc["authors_tks"])
|
||||||
|
# is it English
|
||||||
|
eng = pdf_parser.is_english
|
||||||
|
print("It's English.....", eng)
|
||||||
|
|
||||||
|
res = []
|
||||||
|
# add tables
|
||||||
|
for img, rows in paper["tables"]:
|
||||||
|
bs = 10
|
||||||
|
de = ";" if eng else ";"
|
||||||
|
for i in range(0, len(rows), bs):
|
||||||
|
d = copy.deepcopy(doc)
|
||||||
|
r = de.join(rows[i:i + bs])
|
||||||
|
r = re.sub(r"\t——(来自| in ).*”%s" % de, "", r)
|
||||||
|
tokenize(d, r)
|
||||||
|
d["image"] = img
|
||||||
|
res.append(d)
|
||||||
|
|
||||||
|
if paper["abstract"]:
|
||||||
|
d = copy.deepcopy(doc)
|
||||||
|
txt = pdf_parser.remove_tag(paper["abstract"])
|
||||||
|
d["important_kwd"] = ["abstract", "总结", "概括", "summary", "summarize"]
|
||||||
|
d["important_tks"] = " ".join(d["important_kwd"])
|
||||||
|
d["image"] = pdf_parser.crop(paper["abstract"])
|
||||||
|
tokenize(d, txt, eng)
|
||||||
|
res.append(d)
|
||||||
|
|
||||||
|
readed = [0] * len(paper["lines"])
|
||||||
|
# find colon firstly
|
||||||
|
i = 0
|
||||||
|
while i + 1 < len(paper["lines"]):
|
||||||
|
txt = pdf_parser.remove_tag(paper["lines"][i][0])
|
||||||
|
j = i
|
||||||
|
if txt.strip("\n").strip()[-1] not in "::":
|
||||||
|
i += 1
|
||||||
|
continue
|
||||||
|
i += 1
|
||||||
|
while i < len(paper["lines"]) and not paper["lines"][i][0]:
|
||||||
|
i += 1
|
||||||
|
if i >= len(paper["lines"]): break
|
||||||
|
proj = [paper["lines"][i][0].strip()]
|
||||||
|
i += 1
|
||||||
|
while i < len(paper["lines"]) and paper["lines"][i][0].strip()[0] == proj[-1][0]:
|
||||||
|
proj.append(paper["lines"][i])
|
||||||
|
i += 1
|
||||||
|
for k in range(j, i): readed[k] = True
|
||||||
|
txt = txt[::-1]
|
||||||
|
if eng:
|
||||||
|
r = re.search(r"(.*?) ([\.;?!]|$)", txt)
|
||||||
|
txt = r.group(1)[::-1] if r else txt[::-1]
|
||||||
|
else:
|
||||||
|
r = re.search(r"(.*?) ([。?;!]|$)", txt)
|
||||||
|
txt = r.group(1)[::-1] if r else txt[::-1]
|
||||||
|
for p in proj:
|
||||||
|
d = copy.deepcopy(doc)
|
||||||
|
txt += "\n" + pdf_parser.remove_tag(p)
|
||||||
|
d["image"] = pdf_parser.crop(p)
|
||||||
|
tokenize(d, txt)
|
||||||
|
res.append(d)
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
chunk = []
|
||||||
|
tk_cnt = 0
|
||||||
|
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)
|
||||||
|
d["image"] = pdf_parser.crop(ck)
|
||||||
|
res.append(d)
|
||||||
|
chunk = []
|
||||||
|
tk_cnt = 0
|
||||||
|
|
||||||
|
while i < len(paper["lines"]):
|
||||||
|
if tk_cnt > 128:
|
||||||
|
add_chunk()
|
||||||
|
if readed[i]:
|
||||||
|
i += 1
|
||||||
|
continue
|
||||||
|
readed[i] = True
|
||||||
|
txt, layouts = paper["lines"][i]
|
||||||
|
txt_ = pdf_parser.remove_tag(txt)
|
||||||
|
i += 1
|
||||||
|
cnt = num_tokens_from_string(txt_)
|
||||||
|
if any([
|
||||||
|
layouts.find("title") >= 0 and chunk,
|
||||||
|
cnt + tk_cnt > 128 and tk_cnt > 32,
|
||||||
|
]):
|
||||||
|
add_chunk()
|
||||||
|
chunk = [txt]
|
||||||
|
tk_cnt = cnt
|
||||||
|
else:
|
||||||
|
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")
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import sys
|
||||||
|
|
||||||
|
chunk(sys.argv[1])
|
@ -3,7 +3,7 @@ import re
|
|||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
from pptx import Presentation
|
from pptx import Presentation
|
||||||
|
|
||||||
from rag.app import callback__
|
from rag.app import callback__, tokenize, is_english
|
||||||
from rag.nlp import huqie
|
from rag.nlp import huqie
|
||||||
from rag.parser.pdf_parser import HuParser
|
from rag.parser.pdf_parser import HuParser
|
||||||
|
|
||||||
@ -57,7 +57,7 @@ class Ppt(object):
|
|||||||
assert len(imgs) == len(txts), "Slides text and image do not match: {} vs. {}".format(len(imgs), len(txts))
|
assert len(imgs) == len(txts), "Slides text and image do not match: {} vs. {}".format(len(imgs), len(txts))
|
||||||
callback__((min(to_page, self.total_page) - from_page) / self.total_page,
|
callback__((min(to_page, self.total_page) - from_page) / self.total_page,
|
||||||
"Page {}~{}: Image extraction finished".format(from_page, min(to_page, self.total_page)), callback)
|
"Page {}~{}: Image extraction finished".format(from_page, min(to_page, self.total_page)), callback)
|
||||||
|
self.is_english = is_english(txts)
|
||||||
return [(txts[i], imgs[i]) for i in range(len(txts))]
|
return [(txts[i], imgs[i]) for i in range(len(txts))]
|
||||||
|
|
||||||
|
|
||||||
@ -103,19 +103,19 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
|
|||||||
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
||||||
res = []
|
res = []
|
||||||
if re.search(r"\.pptx?$", filename, re.IGNORECASE):
|
if re.search(r"\.pptx?$", filename, re.IGNORECASE):
|
||||||
for txt,img in Ppt()(filename if not binary else binary, from_page, to_page, callback):
|
ppt_parser = Ppt()
|
||||||
|
for txt,img in ppt_parser(filename if not binary else binary, from_page, to_page, callback):
|
||||||
d = copy.deepcopy(doc)
|
d = copy.deepcopy(doc)
|
||||||
d["content_ltks"] = huqie.qie(txt)
|
|
||||||
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
|
|
||||||
d["image"] = img
|
d["image"] = img
|
||||||
|
tokenize(d, txt, ppt_parser.is_english)
|
||||||
res.append(d)
|
res.append(d)
|
||||||
return res
|
return res
|
||||||
if re.search(r"\.pdf$", filename, re.IGNORECASE):
|
if re.search(r"\.pdf$", filename, re.IGNORECASE):
|
||||||
for txt,img in Pdf()(filename if not binary else binary, from_page=from_page, to_page=to_page, callback=callback):
|
pdf_parser = Pdf()
|
||||||
|
for txt,img in pdf_parser(filename if not binary else binary, from_page=from_page, to_page=to_page, callback=callback):
|
||||||
d = copy.deepcopy(doc)
|
d = copy.deepcopy(doc)
|
||||||
d["content_ltks"] = huqie.qie(txt)
|
|
||||||
d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
|
|
||||||
d["image"] = img
|
d["image"] = img
|
||||||
|
tokenize(d, txt, pdf_parser.is_english)
|
||||||
res.append(d)
|
res.append(d)
|
||||||
return res
|
return res
|
||||||
callback__(-1, "This kind of presentation document did not support yet!", callback)
|
callback__(-1, "This kind of presentation document did not support yet!", callback)
|
||||||
|
@ -2,3 +2,6 @@ from . import search
|
|||||||
from rag.utils import ELASTICSEARCH
|
from rag.utils import ELASTICSEARCH
|
||||||
|
|
||||||
retrievaler = search.Dealer(ELASTICSEARCH)
|
retrievaler = search.Dealer(ELASTICSEARCH)
|
||||||
|
|
||||||
|
from nltk.stem import PorterStemmer
|
||||||
|
stemmer = PorterStemmer()
|
||||||
|
@ -1,4 +1,6 @@
|
|||||||
# -*- coding: utf-8 -*-
|
# -*- coding: utf-8 -*-
|
||||||
|
import random
|
||||||
|
|
||||||
import fitz
|
import fitz
|
||||||
import xgboost as xgb
|
import xgboost as xgb
|
||||||
from io import BytesIO
|
from io import BytesIO
|
||||||
@ -14,6 +16,7 @@ from copy import deepcopy
|
|||||||
from rag.cv.table_recognize import TableTransformer
|
from rag.cv.table_recognize import TableTransformer
|
||||||
from rag.cv.ppdetection import PPDet
|
from rag.cv.ppdetection import PPDet
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
logging.getLogger("pdfminer").setLevel(logging.WARNING)
|
logging.getLogger("pdfminer").setLevel(logging.WARNING)
|
||||||
|
|
||||||
|
|
||||||
@ -22,8 +25,8 @@ class HuParser:
|
|||||||
from paddleocr import PaddleOCR
|
from paddleocr import PaddleOCR
|
||||||
logging.getLogger("ppocr").setLevel(logging.ERROR)
|
logging.getLogger("ppocr").setLevel(logging.ERROR)
|
||||||
self.ocr = PaddleOCR(use_angle_cls=False, lang="ch")
|
self.ocr = PaddleOCR(use_angle_cls=False, lang="ch")
|
||||||
self.layouter = PPDet()
|
self.layouter = PPDet("/data/newpeak/medical-gpt/res/ppdet")
|
||||||
self.tbl_det = TableTransformer()
|
self.tbl_det = PPDet("/data/newpeak/medical-gpt/res/ppdet.tbl")
|
||||||
|
|
||||||
self.updown_cnt_mdl = xgb.Booster()
|
self.updown_cnt_mdl = xgb.Booster()
|
||||||
if torch.cuda.is_available():
|
if torch.cuda.is_available():
|
||||||
@ -141,6 +144,21 @@ class HuParser:
|
|||||||
arr[j + 1] = deepcopy(tmp)
|
arr[j + 1] = deepcopy(tmp)
|
||||||
return arr
|
return arr
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def sort_X_by_page(arr, threashold):
|
||||||
|
# sort using y1 first and then x1
|
||||||
|
arr = sorted(arr, key=lambda r: (r["page_number"], r["x0"], r["top"]))
|
||||||
|
for i in range(len(arr) - 1):
|
||||||
|
for j in range(i, -1, -1):
|
||||||
|
# restore the order using th
|
||||||
|
if abs(arr[j + 1]["x0"] - arr[j]["x0"]) < threashold \
|
||||||
|
and arr[j + 1]["top"] < arr[j]["top"]\
|
||||||
|
and arr[j + 1]["page_number"] == arr[j]["page_number"]:
|
||||||
|
tmp = arr[j]
|
||||||
|
arr[j] = arr[j + 1]
|
||||||
|
arr[j + 1] = tmp
|
||||||
|
return arr
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def sort_R_firstly(arr, thr=0):
|
def sort_R_firstly(arr, thr=0):
|
||||||
# sort using y1 first and then x1
|
# sort using y1 first and then x1
|
||||||
@ -326,7 +344,7 @@ class HuParser:
|
|||||||
return layouts
|
return layouts
|
||||||
|
|
||||||
def __table_paddle(self, images):
|
def __table_paddle(self, images):
|
||||||
tbls = self.tbl_det([img for img in images], threshold=0.5)
|
tbls = self.tbl_det([np.array(img) for img in images], thr=0.5)
|
||||||
res = []
|
res = []
|
||||||
# align left&right for rows, align top&bottom for columns
|
# align left&right for rows, align top&bottom for columns
|
||||||
for tbl in tbls:
|
for tbl in tbls:
|
||||||
@ -482,9 +500,12 @@ class HuParser:
|
|||||||
continue
|
continue
|
||||||
ch = c["bottom"] - c["top"]
|
ch = c["bottom"] - c["top"]
|
||||||
bh = bxs[ii]["bottom"] - bxs[ii]["top"]
|
bh = bxs[ii]["bottom"] - bxs[ii]["top"]
|
||||||
if abs(ch - bh) / max(ch, bh) >= 0.7:
|
if abs(ch - bh) / max(ch, bh) >= 0.7 and c["text"] != ' ':
|
||||||
self.lefted_chars.append(c)
|
self.lefted_chars.append(c)
|
||||||
continue
|
continue
|
||||||
|
if c["text"] == " " and bxs[ii]["text"]:
|
||||||
|
if re.match(r"[0-9a-zA-Z,.?;:!%%]", bxs[ii]["text"][-1]): bxs[ii]["text"] += " "
|
||||||
|
else:
|
||||||
bxs[ii]["text"] += c["text"]
|
bxs[ii]["text"] += c["text"]
|
||||||
|
|
||||||
for b in bxs:
|
for b in bxs:
|
||||||
@ -629,7 +650,7 @@ class HuParser:
|
|||||||
i += 1
|
i += 1
|
||||||
self.boxes = bxs
|
self.boxes = bxs
|
||||||
|
|
||||||
def _concat_downward(self):
|
def _concat_downward(self, concat_between_pages=True):
|
||||||
# count boxes in the same row as a feature
|
# count boxes in the same row as a feature
|
||||||
for i in range(len(self.boxes)):
|
for i in range(len(self.boxes)):
|
||||||
mh = self.mean_height[self.boxes[i]["page_number"] - 1]
|
mh = self.mean_height[self.boxes[i]["page_number"] - 1]
|
||||||
@ -665,6 +686,8 @@ class HuParser:
|
|||||||
if not smpg and ydis > mh * 16:
|
if not smpg and ydis > mh * 16:
|
||||||
break
|
break
|
||||||
down = boxes[i]
|
down = boxes[i]
|
||||||
|
if not concat_between_pages and down["page_number"] > up["page_number"]:
|
||||||
|
break
|
||||||
|
|
||||||
if up.get("R", "") != down.get(
|
if up.get("R", "") != down.get(
|
||||||
"R", "") and up["text"][-1] != ",":
|
"R", "") and up["text"][-1] != ",":
|
||||||
@ -735,43 +758,29 @@ class HuParser:
|
|||||||
|
|
||||||
self.boxes = self.sort_Y_firstly(boxes, 0)
|
self.boxes = self.sort_Y_firstly(boxes, 0)
|
||||||
|
|
||||||
def __filter_forpages(self):
|
def _filter_forpages(self):
|
||||||
if not self.boxes:
|
if not self.boxes:
|
||||||
return
|
return
|
||||||
to = min(7, len(self.page_images) // 5)
|
i = 0
|
||||||
pg_hits = [0 for _ in range(to)]
|
while i < len(self.boxes):
|
||||||
|
if not re.match(r"(contents|目录|目次|table of contents)$", re.sub(r"( | |\u3000)+", "", self.boxes[i]["text"].lower())):
|
||||||
def possible(c):
|
i += 1
|
||||||
if c.get("layout_type", "") == "reference":
|
continue
|
||||||
return True
|
eng = re.match(r"[0-9a-zA-Z :'.-]{5,}", self.boxes[i]["text"].strip())
|
||||||
if c["bottom"] - c["top"] >= 2 * \
|
self.boxes.pop(i)
|
||||||
self.mean_height[c["page_number"] - 1]:
|
if i >= len(self.boxes): break
|
||||||
return False
|
prefix = self.boxes[i]["text"].strip()[:3] if not eng else " ".join(self.boxes[i]["text"].strip().split(" ")[:2])
|
||||||
if c["text"].find("....") >= 0 \
|
while not prefix:
|
||||||
or (c["x1"] - c["x0"] > 250 and re.search(r"[0-9]+$",
|
self.boxes.pop(i)
|
||||||
c["text"].strip())):
|
if i >= len(self.boxes): break
|
||||||
return True
|
prefix = self.boxes[i]["text"].strip()[:3] if not eng else " ".join(self.boxes[i]["text"].strip().split(" ")[:2])
|
||||||
return self.is_caption(c) and re.search(
|
self.boxes.pop(i)
|
||||||
r"[0-9]+$", c["text"].strip())
|
if i >= len(self.boxes) or not prefix: break
|
||||||
|
for j in range(i, min(i + 128, len(self.boxes))):
|
||||||
for c in self.boxes:
|
if not re.match(prefix, self.boxes[j]["text"]):
|
||||||
if c["page_number"] >= to:
|
continue
|
||||||
|
for k in range(i, j): self.boxes.pop(i)
|
||||||
break
|
break
|
||||||
if possible(c):
|
|
||||||
pg_hits[c["page_number"] - 1] += 1
|
|
||||||
|
|
||||||
st, ed = -1, -1
|
|
||||||
for i in range(len(self.boxes)):
|
|
||||||
c = self.boxes[i]
|
|
||||||
if c["page_number"] >= to:
|
|
||||||
break
|
|
||||||
if pg_hits[c["page_number"] - 1] >= 3 and possible(c):
|
|
||||||
if st < 0:
|
|
||||||
st = i
|
|
||||||
else:
|
|
||||||
ed = i
|
|
||||||
for _ in range(st, ed + 1):
|
|
||||||
self.boxes.pop(st)
|
|
||||||
|
|
||||||
def _blockType(self, b):
|
def _blockType(self, b):
|
||||||
patt = [
|
patt = [
|
||||||
@ -1153,6 +1162,7 @@ class HuParser:
|
|||||||
headers = {}
|
headers = {}
|
||||||
hdrset = set()
|
hdrset = set()
|
||||||
lst_hdr = []
|
lst_hdr = []
|
||||||
|
de = "的" if not self.is_english else " for "
|
||||||
for r in sorted(list(hdr_rowno)):
|
for r in sorted(list(hdr_rowno)):
|
||||||
headers[r] = ["" for _ in range(clmno)]
|
headers[r] = ["" for _ in range(clmno)]
|
||||||
for i in range(clmno):
|
for i in range(clmno):
|
||||||
@ -1184,11 +1194,11 @@ class HuParser:
|
|||||||
if headers[j][k].find(headers[j - 1][k]) >= 0:
|
if headers[j][k].find(headers[j - 1][k]) >= 0:
|
||||||
continue
|
continue
|
||||||
if len(headers[j][k]) > len(headers[j - 1][k]):
|
if len(headers[j][k]) > len(headers[j - 1][k]):
|
||||||
headers[j][k] += ("的" if headers[j][k]
|
headers[j][k] += (de if headers[j][k]
|
||||||
else "") + headers[j - 1][k]
|
else "") + headers[j - 1][k]
|
||||||
else:
|
else:
|
||||||
headers[j][k] = headers[j - 1][k] \
|
headers[j][k] = headers[j - 1][k] \
|
||||||
+ ("的" if headers[j - 1][k] else "") \
|
+ (de if headers[j - 1][k] else "") \
|
||||||
+ headers[j][k]
|
+ headers[j][k]
|
||||||
|
|
||||||
logging.debug(
|
logging.debug(
|
||||||
@ -1241,7 +1251,11 @@ class HuParser:
|
|||||||
row_txt.append("; ".join(rtxt))
|
row_txt.append("; ".join(rtxt))
|
||||||
|
|
||||||
if cap:
|
if cap:
|
||||||
row_txt = [t + f"\t——来自“{cap}”" for t in row_txt]
|
if self.is_english:
|
||||||
|
from_ = " in "
|
||||||
|
else:
|
||||||
|
from_ = "来自"
|
||||||
|
row_txt = [t + f"\t——{from_}“{cap}”" for t in row_txt]
|
||||||
return row_txt
|
return row_txt
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@ -1254,7 +1268,7 @@ class HuParser:
|
|||||||
return True
|
return True
|
||||||
return False
|
return False
|
||||||
|
|
||||||
def __extract_table_figure(self, need_image, ZM, return_html):
|
def _extract_table_figure(self, need_image, ZM, return_html):
|
||||||
tables = {}
|
tables = {}
|
||||||
figures = {}
|
figures = {}
|
||||||
# extract figure and table boxes
|
# extract figure and table boxes
|
||||||
@ -1574,8 +1588,14 @@ class HuParser:
|
|||||||
self.page_chars.append([])
|
self.page_chars.append([])
|
||||||
|
|
||||||
logging.info("Images converted.")
|
logging.info("Images converted.")
|
||||||
|
self.is_english = [re.search(r"[a-zA-Z0-9,/¸;:'\[\]\(\)!@#$%^&*\"?<>._-]{30,}", "".join(random.choices([c["text"] for c in self.page_chars[i]], k=100))) for i in range(len(self.page_chars))]
|
||||||
|
if sum([1 if e else 0 for e in self.is_english]) > len(self.page_images) / 2:
|
||||||
|
self.is_english = True
|
||||||
|
else:
|
||||||
|
self.is_english = False
|
||||||
|
|
||||||
for i, img in enumerate(self.page_images):
|
for i, img in enumerate(self.page_images):
|
||||||
chars = self.page_chars[i]
|
chars = self.page_chars[i] if not self.is_english else []
|
||||||
self.mean_height.append(
|
self.mean_height.append(
|
||||||
np.median(sorted([c["height"] for c in chars])) if chars else 0
|
np.median(sorted([c["height"] for c in chars])) if chars else 0
|
||||||
)
|
)
|
||||||
@ -1583,6 +1603,14 @@ class HuParser:
|
|||||||
np.median(sorted([c["width"] for c in chars])) if chars else 8
|
np.median(sorted([c["width"] for c in chars])) if chars else 8
|
||||||
)
|
)
|
||||||
self.page_cum_height.append(img.size[1] / zoomin)
|
self.page_cum_height.append(img.size[1] / zoomin)
|
||||||
|
j = 0
|
||||||
|
while j + 1 < len(chars):
|
||||||
|
if chars[j]["text"] and chars[j + 1]["text"] \
|
||||||
|
and re.match(r"[0-9a-zA-Z,.:;!%]+", chars[j]["text"] + chars[j + 1]["text"]) \
|
||||||
|
and chars[j + 1]["x0"] - chars[j]["x1"] >= min(chars[j + 1]["width"],
|
||||||
|
chars[j]["width"]) / 2:
|
||||||
|
chars[j]["text"] += " "
|
||||||
|
j += 1
|
||||||
# if i > 0:
|
# if i > 0:
|
||||||
# if not chars:
|
# if not chars:
|
||||||
# self.page_cum_height.append(img.size[1] / zoomin)
|
# self.page_cum_height.append(img.size[1] / zoomin)
|
||||||
@ -1591,6 +1619,11 @@ class HuParser:
|
|||||||
# np.max([c["bottom"] for c in chars]))
|
# np.max([c["bottom"] for c in chars]))
|
||||||
self.__ocr_paddle(i + 1, img, chars, zoomin)
|
self.__ocr_paddle(i + 1, img, chars, zoomin)
|
||||||
|
|
||||||
|
if not self.is_english and not all([c for c in self.page_chars]) and self.boxes:
|
||||||
|
self.is_english = re.search(r"[\na-zA-Z0-9,/¸;:'\[\]\(\)!@#$%^&*\"?<>._-]{30,}", "".join([b["text"] for b in random.choices(self.boxes, k=30)]))
|
||||||
|
|
||||||
|
logging.info("Is it English:", self.is_english)
|
||||||
|
|
||||||
self.page_cum_height = np.cumsum(self.page_cum_height)
|
self.page_cum_height = np.cumsum(self.page_cum_height)
|
||||||
assert len(self.page_cum_height) == len(self.page_images) + 1
|
assert len(self.page_cum_height) == len(self.page_images) + 1
|
||||||
|
|
||||||
@ -1600,8 +1633,8 @@ class HuParser:
|
|||||||
self._table_transformer_job(zoomin)
|
self._table_transformer_job(zoomin)
|
||||||
self._text_merge()
|
self._text_merge()
|
||||||
self._concat_downward()
|
self._concat_downward()
|
||||||
self.__filter_forpages()
|
self._filter_forpages()
|
||||||
tbls = self.__extract_table_figure(need_image, zoomin, return_html)
|
tbls = self._extract_table_figure(need_image, zoomin, return_html)
|
||||||
return self.__filterout_scraps(deepcopy(self.boxes), zoomin), tbls
|
return self.__filterout_scraps(deepcopy(self.boxes), zoomin), tbls
|
||||||
|
|
||||||
def remove_tag(self, txt):
|
def remove_tag(self, txt):
|
||||||
|
Loading…
x
Reference in New Issue
Block a user