add use layout or not option (#145)

* add use layout or not option

* trival

rag/app/manual.py

@@ -2,8 +2,8 @@ import copy
 import re
 
 from api.db import ParserType
-from rag.nlp import huqie, tokenize, tokenize_table, add_positions, bullets_category, title_frequency
-from deepdoc.parser import PdfParser
+from rag.nlp import huqie, tokenize, tokenize_table, add_positions, bullets_category, title_frequency, tokenize_chunks
+from deepdoc.parser import PdfParser, PlainParser
 from rag.utils import num_tokens_from_string
 
 
@@ -30,9 +30,7 @@ class Pdf(PdfParser):
         #        print(b)
         print("OCR:", timer()-start)
 
-        def tag(pn, left, right, top, bottom):
-            return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##" \
-                .format(pn, left, right, top, bottom)
+
 
         self._layouts_rec(zoomin)
         callback(0.65, "Layout analysis finished.")
@@ -49,6 +47,8 @@ class Pdf(PdfParser):
         for b in self.boxes:
             b["text"] = re.sub(r"([\t 　]|\u3000){2,}", " ", b["text"].strip())
 
+        return [(b["text"], b.get("layout_no", ""), self.get_position(b, zoomin)) for i, b in enumerate(self.boxes)]
+
         # set pivot using the most frequent type of title,
         # then merge between 2 pivot
         if len(self.boxes)>0 and len(self.outlines)/len(self.boxes) > 0.1:
@@ -103,9 +103,10 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
     pdf_parser = None
 
     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)
+        pdf_parser = Pdf() if kwargs.get("parser_config",{}).get("layout_recognize", True) else PlainParser()
+        sections, tbls = pdf_parser(filename if not binary else binary,
+                               from_page=from_page, to_page=to_page, callback=callback)
+        if sections and len(sections[0])<3: cks = [(t, l, [0]*5) for t, l in sections]
     else: raise NotImplementedError("file type not supported yet(pdf supported)")
     doc = {
         "docnm_kwd": filename
@@ -115,13 +116,60 @@ 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
 
+    # set pivot using the most frequent type of title,
+    # then merge between 2 pivot
+    if len(sections) > 0 and len(pdf_parser.outlines) / len(sections) > 0.1:
+        max_lvl = max([lvl for _, lvl in pdf_parser.outlines])
+        most_level = max(0, max_lvl - 1)
+        levels = []
+        for txt, _, _ in sections:
+            for t, lvl in pdf_parser.outlines:
+                tks = set([t[i] + t[i + 1] for i in range(len(t) - 1)])
+                tks_ = set([txt[i] + txt[i + 1] for i in range(min(len(t), len(txt) - 1))])
+                if len(set(tks & tks_)) / max([len(tks), len(tks_), 1]) > 0.8:
+                    levels.append(lvl)
+                    break
+            else:
+                levels.append(max_lvl + 1)
+    else:
+        bull = bullets_category([txt for txt,_,_ in sections])
+        most_level, levels = title_frequency(bull, [(txt, l) for txt, l, poss in sections])
+
+    assert len(sections) == len(levels)
+    sec_ids = []
+    sid = 0
+    for i, lvl in enumerate(levels):
+        if lvl <= most_level and i > 0 and lvl != levels[i - 1]: sid += 1
+        sec_ids.append(sid)
+        # print(lvl, self.boxes[i]["text"], most_level, sid)
+
+    sections = [(txt, sec_ids[i], poss) for i, (txt, _, poss) in enumerate(sections)]
+    for (img, rows), poss in tbls:
+        sections.append((rows if isinstance(rows, str) else rows[0], -1,
+                         [(p[0] + 1 - from_page, p[1], p[2], p[3], p[4]) for p in poss]))
+
+    def tag(pn, left, right, top, bottom):
+        if pn+left+right+top+bottom == 0:
+            return ""
+        return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##" \
+            .format(pn, left, right, top, bottom)
+
+    chunks = []
+    last_sid = -2
+    tk_cnt = 0
+    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 tk_cnt < 2048 and (sec_id == last_sid or sec_id == -1):
+            if chunks:
+                chunks[-1] += "\n" + txt + poss
+                tk_cnt += num_tokens_from_string(txt)
+                continue
+        chunks.append(txt + poss)
+        tk_cnt = num_tokens_from_string(txt)
+        if sec_id > -1: last_sid = sec_id
+
     res = tokenize_table(tbls, doc, eng)
-    for ck in cks:
-        d = copy.deepcopy(doc)
-        d["image"], poss = pdf_parser.crop(ck, need_position=True)
-        add_positions(d, poss)
-        tokenize(d, pdf_parser.remove_tag(ck), eng)
-        res.append(d)
+    res.extend(tokenize_chunks(chunks, doc, eng, pdf_parser))
     return res