| if not e: | if not e: | ||||
| return get_data_error_result(retmsg="Document not found!") | return get_data_error_result(retmsg="Document not found!") | ||||
| query = { | query = { | ||||
| "doc_ids": [doc_id], "page": page, "size": size, "question": question | |||||
| "doc_ids": [doc_id], "page": page, "size": size, "question": question, "sort": True | |||||
| } | } | ||||
| if "available_int" in req: | if "available_int" in req: | ||||
| query["available_int"] = int(req["available_int"]) | query["available_int"] = int(req["available_int"]) | ||||
| "important_kwd": sres.field[id].get("important_kwd", []), | "important_kwd": sres.field[id].get("important_kwd", []), | ||||
| "img_id": sres.field[id].get("img_id", ""), | "img_id": sres.field[id].get("img_id", ""), | ||||
| "available_int": sres.field[id].get("available_int", 1), | "available_int": sres.field[id].get("available_int", 1), | ||||
| "positions": sres.field[id].get("position_int", "").split("\t") | |||||
| } | } | ||||
| poss = [] | |||||
| for i in range(0, len(d["positions"]), 5): | |||||
| poss.append([float(d["positions"][i]), float(d["positions"][i+1]), float(d["positions"][i+2]), float(d["positions"][i+3]), float(d["positions"][i+4])]) | |||||
| d["positions"] = poss | |||||
| res["chunks"].append(d) | res["chunks"].append(d) | ||||
| return get_json_result(data=res) | return get_json_result(data=res) | ||||
| except Exception as e: | except Exception as e: |
| from api.utils.file_utils import get_project_base_directory | from api.utils.file_utils import get_project_base_directory | ||||
| from api.utils.log_utils import LoggerFactory, getLogger | from api.utils.log_utils import LoggerFactory, getLogger | ||||
| from rag.nlp import search | |||||
| from rag.utils import ELASTICSEARCH | |||||
| # Logger | |||||
| LoggerFactory.set_directory(os.path.join(get_project_base_directory(), "logs", "api")) | |||||
| # {CRITICAL: 50, FATAL:50, ERROR:40, WARNING:30, WARN:30, INFO:20, DEBUG:10, NOTSET:0} | |||||
| LoggerFactory.LEVEL = 10 | |||||
| stat_logger = getLogger("stat") | |||||
| access_logger = getLogger("access") | |||||
| database_logger = getLogger("database") | |||||
| API_VERSION = "v1" | API_VERSION = "v1" | ||||
| RAG_FLOW_SERVICE_NAME = "ragflow" | RAG_FLOW_SERVICE_NAME = "ragflow" | ||||
| PRIVILEGE_COMMAND_WHITELIST = [] | PRIVILEGE_COMMAND_WHITELIST = [] | ||||
| CHECK_NODES_IDENTITY = False | CHECK_NODES_IDENTITY = False | ||||
| from rag.nlp import search | |||||
| from rag.utils import ELASTICSEARCH | |||||
| retrievaler = search.Dealer(ELASTICSEARCH) | retrievaler = search.Dealer(ELASTICSEARCH) | ||||
| # Logger | |||||
| LoggerFactory.set_directory(os.path.join(get_project_base_directory(), "logs", "api")) | |||||
| # {CRITICAL: 50, FATAL:50, ERROR:40, WARNING:30, WARN:30, INFO:20, DEBUG:10, NOTSET:0} | |||||
| LoggerFactory.LEVEL = 10 | |||||
| stat_logger = getLogger("stat") | |||||
| access_logger = getLogger("access") | |||||
| database_logger = getLogger("database") | |||||
| class CustomEnum(Enum): | class CustomEnum(Enum): | ||||
| @classmethod | @classmethod |
| b_["top"] = b["top"] | b_["top"] = b["top"] | ||||
| self.boxes.pop(i) | self.boxes.pop(i) | ||||
| def _extract_table_figure(self, need_image, ZM, return_html): | |||||
| def _extract_table_figure(self, need_image, ZM, return_html, need_position): | |||||
| tables = {} | tables = {} | ||||
| figures = {} | figures = {} | ||||
| # extract figure and table boxes | # extract figure and table boxes | ||||
| self.boxes.pop(i) | self.boxes.pop(i) | ||||
| res = [] | res = [] | ||||
| positions = [] | |||||
| def cropout(bxs, ltype): | |||||
| def cropout(bxs, ltype, poss): | |||||
| nonlocal ZM | nonlocal ZM | ||||
| pn = set([b["page_number"] - 1 for b in bxs]) | pn = set([b["page_number"] - 1 for b in bxs]) | ||||
| if len(pn) < 2: | if len(pn) < 2: | ||||
| "layoutno", ""))) | "layoutno", ""))) | ||||
| left, top, right, bott = b["x0"], b["top"], b["x1"], b["bottom"] | left, top, right, bott = b["x0"], b["top"], b["x1"], b["bottom"] | ||||
| poss.append((pn, left, right, top, bott)) | |||||
| return self.page_images[pn] \ | return self.page_images[pn] \ | ||||
| .crop((left * ZM, top * ZM, | .crop((left * ZM, top * ZM, | ||||
| right * ZM, bott * ZM)) | right * ZM, bott * ZM)) | ||||
| pn[p] = [] | pn[p] = [] | ||||
| pn[p].append(b) | pn[p].append(b) | ||||
| pn = sorted(pn.items(), key=lambda x: x[0]) | pn = sorted(pn.items(), key=lambda x: x[0]) | ||||
| imgs = [cropout(arr, ltype) for p, arr in pn] | |||||
| imgs = [cropout(arr, ltype, poss) for p, arr in pn] | |||||
| pic = Image.new("RGB", | pic = Image.new("RGB", | ||||
| (int(np.max([i.size[0] for i in imgs])), | (int(np.max([i.size[0] for i in imgs])), | ||||
| int(np.sum([m.size[1] for m in imgs]))), | int(np.sum([m.size[1] for m in imgs]))), | ||||
| if not txt: | if not txt: | ||||
| continue | continue | ||||
| poss = [] | |||||
| res.append( | res.append( | ||||
| (cropout( | (cropout( | ||||
| bxs, | bxs, | ||||
| "figure"), | |||||
| "figure", poss), | |||||
| [txt] if not return_html else [f"<p>{txt}</p>"])) | [txt] if not return_html else [f"<p>{txt}</p>"])) | ||||
| positions.append(poss) | |||||
| for k, bxs in tables.items(): | for k, bxs in tables.items(): | ||||
| if not bxs: | if not bxs: | ||||
| continue | continue | ||||
| res.append((cropout(bxs, "table"), | |||||
| bxs = Recognizer.sort_Y_firstly(bxs, np.mean([(b["bottom"]-b["top"])/2 for b in bxs])) | |||||
| poss = [] | |||||
| res.append((cropout(bxs, "table", poss), | |||||
| self.tbl_det.construct_table(bxs, html=return_html, is_english=self.is_english))) | self.tbl_det.construct_table(bxs, html=return_html, is_english=self.is_english))) | ||||
| positions.append(poss) | |||||
| assert len(positions) == len(res) | |||||
| if need_position: return list(zip(res, positions)) | |||||
| return res | return res | ||||
| def proj_match(self, line): | def proj_match(self, line): | ||||
| 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, False) | |||||
| 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): | ||||
| return re.sub(r"@@[\t0-9.-]+?##", "", txt) | return re.sub(r"@@[\t0-9.-]+?##", "", txt) | ||||
| def crop(self, text, ZM=3): | |||||
| def crop(self, text, ZM=3, need_position=False): | |||||
| imgs = [] | imgs = [] | ||||
| poss = [] | poss = [] | ||||
| for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", text): | for tag in re.findall(r"@@[0-9-]+\t[0-9.\t]+##", text): | ||||
| pos = poss[-1] | pos = poss[-1] | ||||
| poss.append(([pos[0][-1]], pos[1], pos[2], min(self.page_images[pos[0][-1]].size[1]/ZM, pos[4]+GAP), min(self.page_images[pos[0][-1]].size[1]/ZM, pos[4]+120))) | poss.append(([pos[0][-1]], pos[1], pos[2], min(self.page_images[pos[0][-1]].size[1]/ZM, pos[4]+GAP), min(self.page_images[pos[0][-1]].size[1]/ZM, pos[4]+120))) | ||||
| positions = [] | |||||
| for ii, (pns, left, right, top, bottom) in enumerate(poss): | for ii, (pns, left, right, top, bottom) in enumerate(poss): | ||||
| right = left + max_width | right = left + max_width | ||||
| bottom *= ZM | bottom *= ZM | ||||
| bottom, self.page_images[pns[0]].size[1]) | bottom, self.page_images[pns[0]].size[1]) | ||||
| )) | )) | ||||
| ) | ) | ||||
| positions.append((pns[0], left, right, top, min( | |||||
| bottom, self.page_images[pns[0]].size[1])/ZM)) | |||||
| bottom -= self.page_images[pns[0]].size[1] | bottom -= self.page_images[pns[0]].size[1] | ||||
| for pn in pns[1:]: | for pn in pns[1:]: | ||||
| imgs.append( | imgs.append( | ||||
| self.page_images[pn].size[1]) | self.page_images[pn].size[1]) | ||||
| )) | )) | ||||
| ) | ) | ||||
| positions.append((pn, left, right, 0, min( | |||||
| bottom, self.page_images[pn].size[1]) / ZM)) | |||||
| bottom -= self.page_images[pn].size[1] | bottom -= self.page_images[pn].size[1] | ||||
| if not imgs: | if not imgs: | ||||
| if need_position: return None, None | |||||
| return | return | ||||
| height = 0 | height = 0 | ||||
| for img in imgs: | for img in imgs: | ||||
| img = Image.alpha_composite(img, overlay).convert("RGB") | img = Image.alpha_composite(img, overlay).convert("RGB") | ||||
| pic.paste(img, (0, int(height))) | pic.paste(img, (0, int(height))) | ||||
| height += img.size[1] + GAP | height += img.size[1] + GAP | ||||
| if need_position: | |||||
| return pic, positions | |||||
| return pic | return pic | ||||
| return | return | ||||
| min_dis, min_i = 1000000, None | min_dis, min_i = 1000000, None | ||||
| for i,b in enumerate(boxes): | for i,b in enumerate(boxes): | ||||
| if box.get("layoutno", "0") != b.get("layoutno", "0"): continue | |||||
| dis = min(abs(box["x0"] - b["x0"]), abs(box["x1"] - b["x1"]), abs(box["x0"]+box["x1"] - b["x1"] - b["x0"])/2) | dis = min(abs(box["x0"] - b["x0"]), abs(box["x1"] - b["x1"]), abs(box["x0"]+box["x1"] - b["x1"] - b["x0"])/2) | ||||
| if dis < min_dis: | if dis < min_dis: | ||||
| min_i = i | min_i = i |
| import copy | import copy | ||||
| import re | import re | ||||
| from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, \ | from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, \ | ||||
| hierarchical_merge, make_colon_as_title, naive_merge, random_choices, tokenize_table | |||||
| hierarchical_merge, make_colon_as_title, naive_merge, random_choices, tokenize_table, add_positions | |||||
| from rag.nlp import huqie | from rag.nlp import huqie | ||||
| from deepdoc.parser import PdfParser, DocxParser | from deepdoc.parser import PdfParser, DocxParser | ||||
| class Pdf(PdfParser): | class Pdf(PdfParser): | ||||
| def __call__(self, filename, binary=None, from_page=0, | def __call__(self, filename, binary=None, from_page=0, | ||||
| to_page=100000, zoomin=3, callback=None): | to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__( | self.__images__( | ||||
| filename if not binary else binary, | filename if not binary else binary, | ||||
| zoomin, | zoomin, | ||||
| self._filter_forpages() | self._filter_forpages() | ||||
| self._merge_with_same_bullet() | self._merge_with_same_bullet() | ||||
| callback(0.75, "Text merging finished.") | callback(0.75, "Text merging finished.") | ||||
| tbls = self._extract_table_figure(True, zoomin, False) | |||||
| tbls = self._extract_table_figure(True, zoomin, False, True) | |||||
| callback(0.8, "Text extraction finished") | callback(0.8, "Text extraction finished") | ||||
| return [(b["text"] + self._line_tag(b, zoomin), b.get("layoutno","")) for b in self.boxes], tbls | |||||
| return [(b["text"] + self._line_tag(b, zoomin), b.get("layoutno","")) for b in self.boxes], tbls, tbl_poss | |||||
| def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", callback=None, **kwargs): | def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", callback=None, **kwargs): | ||||
| callback(0.8, "Finish parsing.") | callback(0.8, "Finish parsing.") | ||||
| elif re.search(r"\.pdf$", filename, re.IGNORECASE): | elif re.search(r"\.pdf$", filename, re.IGNORECASE): | ||||
| pdf_parser = Pdf() | pdf_parser = Pdf() | ||||
| sections,tbls = pdf_parser(filename if not binary else binary, | |||||
| sections, tbls = pdf_parser(filename if not binary else binary, | |||||
| from_page=from_page, to_page=to_page, callback=callback) | from_page=from_page, to_page=to_page, callback=callback) | ||||
| elif re.search(r"\.txt$", filename, re.IGNORECASE): | elif re.search(r"\.txt$", filename, re.IGNORECASE): | ||||
| callback(0.1, "Start to parse.") | callback(0.1, "Start to parse.") | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| ck = "\n".join(ck) | ck = "\n".join(ck) | ||||
| if pdf_parser: | if pdf_parser: | ||||
| d["image"] = pdf_parser.crop(ck) | |||||
| d["image"], poss = pdf_parser.crop(ck, need_position=True) | |||||
| add_positions(d, poss) | |||||
| ck = pdf_parser.remove_tag(ck) | ck = pdf_parser.remove_tag(ck) | ||||
| tokenize(d, ck, eng) | tokenize(d, ck, eng) | ||||
| res.append(d) | res.append(d) |
| from io import BytesIO | from io import BytesIO | ||||
| from docx import Document | from docx import Document | ||||
| from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, hierarchical_merge, \ | from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, hierarchical_merge, \ | ||||
| make_colon_as_title | |||||
| make_colon_as_title, add_positions | |||||
| from rag.nlp import huqie | from rag.nlp import huqie | ||||
| from deepdoc.parser import PdfParser, DocxParser | from deepdoc.parser import PdfParser, DocxParser | ||||
| from rag.settings import cron_logger | from rag.settings import cron_logger | ||||
| class Pdf(PdfParser): | class Pdf(PdfParser): | ||||
| def __call__(self, filename, binary=None, from_page=0, | def __call__(self, filename, binary=None, from_page=0, | ||||
| to_page=100000, zoomin=3, callback=None): | to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__( | self.__images__( | ||||
| filename if not binary else binary, | filename if not binary else binary, | ||||
| zoomin, | zoomin, | ||||
| ck = "\n".join(ck) | ck = "\n".join(ck) | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| if pdf_parser: | if pdf_parser: | ||||
| d["image"] = pdf_parser.crop(ck) | |||||
| d["image"], poss = pdf_parser.crop(ck, need_position=True) | |||||
| add_positions(d, poss) | |||||
| ck = pdf_parser.remove_tag(ck) | ck = pdf_parser.remove_tag(ck) | ||||
| tokenize(d, ck, eng) | tokenize(d, ck, eng) | ||||
| res.append(d) | res.append(d) |
| import re | import re | ||||
| from api.db import ParserType | from api.db import ParserType | ||||
| from rag.nlp import huqie, tokenize, tokenize_table | |||||
| from rag.nlp import huqie, tokenize, tokenize_table, add_positions | |||||
| from deepdoc.parser import PdfParser | from deepdoc.parser import PdfParser | ||||
| from rag.utils import num_tokens_from_string | from rag.utils import num_tokens_from_string | ||||
| def __call__(self, filename, binary=None, from_page=0, | def __call__(self, filename, binary=None, from_page=0, | ||||
| to_page=100000, zoomin=3, callback=None): | to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__( | self.__images__( | ||||
| filename if not binary else binary, | filename if not binary else binary, | ||||
| zoomin, | zoomin, | ||||
| self._concat_downward(concat_between_pages=False) | self._concat_downward(concat_between_pages=False) | ||||
| self._filter_forpages() | self._filter_forpages() | ||||
| callback(0.77, "Text merging finished") | callback(0.77, "Text merging finished") | ||||
| tbls = self._extract_table_figure(True, zoomin, False) | |||||
| tbls = self._extract_table_figure(True, zoomin, False, True) | |||||
| # clean mess | # clean mess | ||||
| for b in self.boxes: | for b in self.boxes: | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| ck = "\n".join(chunk) | ck = "\n".join(chunk) | ||||
| tokenize(d, pdf_parser.remove_tag(ck), pdf_parser.is_english) | tokenize(d, pdf_parser.remove_tag(ck), pdf_parser.is_english) | ||||
| d["image"] = pdf_parser.crop(ck) | |||||
| d["image"], poss = pdf_parser.crop(ck, need_position=True) | |||||
| add_positions(d, poss) | |||||
| res.append(d) | res.append(d) | ||||
| chunk = [] | chunk = [] | ||||
| tk_cnt = 0 | tk_cnt = 0 |
| import copy | import copy | ||||
| import re | import re | ||||
| from rag.app import laws | from rag.app import laws | ||||
| from rag.nlp import huqie, is_english, tokenize, naive_merge, tokenize_table | |||||
| from rag.nlp import huqie, is_english, tokenize, naive_merge, tokenize_table, add_positions | |||||
| from deepdoc.parser import PdfParser | from deepdoc.parser import PdfParser | ||||
| from rag.settings import cron_logger | from rag.settings import cron_logger | ||||
| class Pdf(PdfParser): | class Pdf(PdfParser): | ||||
| def __call__(self, filename, binary=None, from_page=0, | def __call__(self, filename, binary=None, from_page=0, | ||||
| to_page=100000, zoomin=3, callback=None): | to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__( | self.__images__( | ||||
| filename if not binary else binary, | filename if not binary else binary, | ||||
| zoomin, | zoomin, | ||||
| self._concat_downward(concat_between_pages=False) | self._concat_downward(concat_between_pages=False) | ||||
| self._filter_forpages() | self._filter_forpages() | ||||
| callback(0.77, "Text merging finished") | callback(0.77, "Text merging finished") | ||||
| tbls = self._extract_table_figure(True, zoomin, False) | |||||
| tbls = self._extract_table_figure(True, zoomin, False, True) | |||||
| cron_logger.info("paddle layouts:".format((timer() - start) / (self.total_page + 0.1))) | cron_logger.info("paddle layouts:".format((timer() - start) / (self.total_page + 0.1))) | ||||
| #self._naive_vertical_merge() | #self._naive_vertical_merge() | ||||
| # wrap up to es documents | # wrap up to es documents | ||||
| for ck in cks: | for ck in cks: | ||||
| if len(ck.strip()) == 0:continue | |||||
| print("--", ck) | print("--", ck) | ||||
| if not ck:continue | |||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| if pdf_parser: | if pdf_parser: | ||||
| d["image"] = pdf_parser.crop(ck) | |||||
| d["image"], poss = pdf_parser.crop(ck, need_position=True) | |||||
| add_positions(d, poss) | |||||
| ck = pdf_parser.remove_tag(ck) | ck = pdf_parser.remove_tag(ck) | ||||
| tokenize(d, ck, eng) | tokenize(d, ck, eng) | ||||
| res.append(d) | res.append(d) |
| from collections import Counter | from collections import Counter | ||||
| from api.db import ParserType | from api.db import ParserType | ||||
| from rag.nlp import huqie, tokenize, tokenize_table | |||||
| from rag.nlp import huqie, tokenize, tokenize_table, add_positions | |||||
| from deepdoc.parser import PdfParser | from deepdoc.parser import PdfParser | ||||
| import numpy as np | import numpy as np | ||||
| from rag.utils import num_tokens_from_string | from rag.utils import num_tokens_from_string | ||||
| def __call__(self, filename, binary=None, from_page=0, | def __call__(self, filename, binary=None, from_page=0, | ||||
| to_page=100000, zoomin=3, callback=None): | to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__( | self.__images__( | ||||
| filename if not binary else binary, | filename if not binary else binary, | ||||
| zoomin, | zoomin, | ||||
| self._concat_downward(concat_between_pages=False) | self._concat_downward(concat_between_pages=False) | ||||
| self._filter_forpages() | self._filter_forpages() | ||||
| callback(0.75, "Text merging finished.") | callback(0.75, "Text merging finished.") | ||||
| tbls = self._extract_table_figure(True, zoomin, False) | |||||
| tbls = self._extract_table_figure(True, zoomin, False, True) | |||||
| # clean mess | # clean mess | ||||
| if column_width < self.page_images[0].size[0] / zoomin / 2: | if column_width < self.page_images[0].size[0] / zoomin / 2: | ||||
| txt = pdf_parser.remove_tag(paper["abstract"]) | txt = pdf_parser.remove_tag(paper["abstract"]) | ||||
| d["important_kwd"] = ["abstract", "总结", "概括", "summary", "summarize"] | d["important_kwd"] = ["abstract", "总结", "概括", "summary", "summarize"] | ||||
| d["important_tks"] = " ".join(d["important_kwd"]) | d["important_tks"] = " ".join(d["important_kwd"]) | ||||
| d["image"] = pdf_parser.crop(paper["abstract"]) | |||||
| d["image"], poss = pdf_parser.crop(paper["abstract"], need_position=True) | |||||
| add_positions(d, poss) | |||||
| tokenize(d, txt, eng) | tokenize(d, txt, eng) | ||||
| res.append(d) | res.append(d) | ||||
| for p in proj: | for p in proj: | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| txt += "\n" + pdf_parser.remove_tag(p) | txt += "\n" + pdf_parser.remove_tag(p) | ||||
| d["image"] = pdf_parser.crop(p) | |||||
| tokenize(d, txt) | |||||
| d["image"], poss = pdf_parser.crop(p, need_position=True) | |||||
| add_positions(d, poss) | |||||
| tokenize(d, txt, eng) | |||||
| res.append(d) | res.append(d) | ||||
| i = 0 | i = 0 | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| ck = "\n".join(chunk) | ck = "\n".join(chunk) | ||||
| tokenize(d, pdf_parser.remove_tag(ck), pdf_parser.is_english) | tokenize(d, pdf_parser.remove_tag(ck), pdf_parser.is_english) | ||||
| d["image"] = pdf_parser.crop(ck) | |||||
| d["image"], poss = pdf_parser.crop(ck, need_position=True) | |||||
| add_positions(d, poss) | |||||
| res.append(d) | res.append(d) | ||||
| chunk = [] | chunk = [] | ||||
| tk_cnt = 0 | tk_cnt = 0 |
| return False | return False | ||||
| def __call__(self, filename, binary=None, from_page=0, to_page=100000, zoomin=3, callback=None): | def __call__(self, filename, binary=None, from_page=0, to_page=100000, zoomin=3, callback=None): | ||||
| callback(msg="OCR is running...") | |||||
| self.__images__(filename if not binary else binary, zoomin, from_page, to_page) | self.__images__(filename if not binary else binary, zoomin, from_page, to_page) | ||||
| callback(0.8, "Page {}~{}: OCR finished".format(from_page, min(to_page, self.total_page))) | callback(0.8, "Page {}~{}: OCR finished".format(from_page, min(to_page, self.total_page))) | ||||
| assert len(self.boxes) == len(self.page_images), "{} vs. {}".format(len(self.boxes), len(self.page_images)) | assert len(self.boxes) == len(self.page_images), "{} vs. {}".format(len(self.boxes), len(self.page_images)) | ||||
| return res | return res | ||||
| elif re.search(r"\.pdf$", filename, re.IGNORECASE): | elif re.search(r"\.pdf$", filename, re.IGNORECASE): | ||||
| pdf_parser = Pdf() | 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): | |||||
| for pn, (txt,img) in enumerate(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["image"] = img | d["image"] = img | ||||
| d["page_num_obj"] = [pn+1] | |||||
| tokenize(d, txt, pdf_parser.is_english) | tokenize(d, txt, pdf_parser.is_english) | ||||
| res.append(d) | res.append(d) | ||||
| return res | return res |
| def tokenize_table(tbls, doc, eng, batch_size=10): | def tokenize_table(tbls, doc, eng, batch_size=10): | ||||
| res = [] | res = [] | ||||
| # add tables | # add tables | ||||
| for img, rows in tbls: | |||||
| for (img, rows), poss in tbls: | |||||
| if not rows:continue | |||||
| if isinstance(rows, str): | |||||
| d = copy.deepcopy(doc) | |||||
| r = re.sub(r"<[^<>]{,12}>", "", rows) | |||||
| tokenize(d, r, eng) | |||||
| d["content_with_weight"] = rows | |||||
| d["image"] = img | |||||
| add_positions(d, poss) | |||||
| res.append(d) | |||||
| continue | |||||
| de = "; " if eng else "; " | de = "; " if eng else "; " | ||||
| for i in range(0, len(rows), batch_size): | for i in range(0, len(rows), batch_size): | ||||
| d = copy.deepcopy(doc) | d = copy.deepcopy(doc) | ||||
| r = de.join(rows[i:i + batch_size]) | r = de.join(rows[i:i + batch_size]) | ||||
| tokenize(d, r, eng) | tokenize(d, r, eng) | ||||
| d["image"] = img | d["image"] = img | ||||
| add_positions(d, poss) | |||||
| res.append(d) | res.append(d) | ||||
| return res | return res | ||||
| def add_positions(d, poss): | |||||
| if not poss:return | |||||
| d["page_num_int"] = [] | |||||
| d["position_int"] = [] | |||||
| d["top_int"] = [] | |||||
| for pn, left, right, top, bottom in poss: | |||||
| d["page_num_int"].append(pn+1) | |||||
| d["top_int"].append(top) | |||||
| d["position_int"].append((pn+1, left, right, top, bottom)) | |||||
| def remove_contents_table(sections, eng=False): | def remove_contents_table(sections, eng=False): | ||||
| i = 0 | i = 0 | ||||
| while i < len(sections): | while i < len(sections): |
| pg = int(req.get("page", 1)) - 1 | pg = int(req.get("page", 1)) - 1 | ||||
| ps = int(req.get("size", 1000)) | ps = int(req.get("size", 1000)) | ||||
| src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id", "img_id", | src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id", "img_id", | ||||
| "image_id", "doc_id", "q_512_vec", "q_768_vec", | |||||
| "image_id", "doc_id", "q_512_vec", "q_768_vec", "position_int", | |||||
| "q_1024_vec", "q_1536_vec", "available_int", "content_with_weight"]) | "q_1024_vec", "q_1536_vec", "available_int", "content_with_weight"]) | ||||
| s = s.query(bqry)[pg * ps:(pg + 1) * ps] | s = s.query(bqry)[pg * ps:(pg + 1) * ps] | ||||
| s = s.highlight("content_ltks") | s = s.highlight("content_ltks") | ||||
| s = s.highlight("title_ltks") | s = s.highlight("title_ltks") | ||||
| if not qst: | if not qst: | ||||
| s = s.sort( | |||||
| {"create_time": {"order": "desc", "unmapped_type": "date"}}, | |||||
| {"create_timestamp_flt": {"order": "desc", "unmapped_type": "float"}} | |||||
| ) | |||||
| if not req.get("sort"): | |||||
| s = s.sort( | |||||
| {"create_time": {"order": "desc", "unmapped_type": "date"}}, | |||||
| {"create_timestamp_flt": {"order": "desc", "unmapped_type": "float"}} | |||||
| ) | |||||
| else: | |||||
| s = s.sort( | |||||
| {"page_num_int": {"order": "asc", "unmapped_type": "float"}}, | |||||
| {"top_int": {"order": "asc", "unmapped_type": "float"}}, | |||||
| {"create_time": {"order": "desc", "unmapped_type": "date"}}, | |||||
| {"create_timestamp_flt": {"order": "desc", "unmapped_type": "float"}} | |||||
| ) | |||||
| if qst: | if qst: | ||||
| s = s.highlight_options( | s = s.highlight_options( | ||||
| m = {n: d.get(n) for n in flds if d.get(n) is not None} | m = {n: d.get(n) for n in flds if d.get(n) is not None} | ||||
| for n, v in m.items(): | for n, v in m.items(): | ||||
| if isinstance(v, type([])): | if isinstance(v, type([])): | ||||
| m[n] = "\t".join([str(vv) for vv in v]) | |||||
| m[n] = "\t".join([str(vv) if not isinstance(vv, list) else "\t".join([str(vvv) for vvv in vv]) for vv in v]) | |||||
| continue | continue | ||||
| if not isinstance(v, type("")): | if not isinstance(v, type("")): | ||||
| m[n] = str(m[n]) | m[n] = str(m[n]) |
| BATCH_SIZE = 64 | BATCH_SIZE = 64 | ||||
| FACTORY = { | FACTORY = { | ||||
| "general": naive, | |||||
| ParserType.NAIVE.value: naive, | ParserType.NAIVE.value: naive, | ||||
| ParserType.PAPER.value: paper, | ParserType.PAPER.value: paper, | ||||
| ParserType.BOOK.value: book, | ParserType.BOOK.value: book, | ||||
| es_r = ELASTICSEARCH.bulk(cks, search.index_name(r["tenant_id"])) | es_r = ELASTICSEARCH.bulk(cks, search.index_name(r["tenant_id"])) | ||||
| if es_r: | if es_r: | ||||
| callback(-1, "Index failure!") | callback(-1, "Index failure!") | ||||
| ELASTICSEARCH.deleteByQuery( | |||||
| Q("match", doc_id=r["doc_id"]), idxnm=search.index_name(r["tenant_id"])) | |||||
| cron_logger.error(str(es_r)) | cron_logger.error(str(es_r)) | ||||
| else: | else: | ||||
| if TaskService.do_cancel(r["id"]): | if TaskService.do_cancel(r["id"]): |