add paper & manual parser (#46)

rag/app/__init__.py

@@ -1,5 +1,9 @@
 import re
 
+from nltk import word_tokenize
+
+from rag.nlp import stemmer, huqie
+
 
 def callback__(progress, msg, func):
     if not func :return
@@ -46,3 +50,21 @@ def bullets_category(sections):
         res = i
         maxium = h
     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"])

rag/app/laws.py

@@ -3,12 +3,13 @@ import re
 from io import BytesIO
 from docx import Document
 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.parser.docx_parser import HuDocxParser
 from rag.parser.pdf_parser import HuParser
 
 
-class Docx(object):
+class Docx(HuDocxParser):
     def __init__(self):
         pass
 
@@ -42,14 +43,7 @@ class Pdf(HuParser):
         print("paddle layouts:", timer()-start)
         bxs = self.sort_Y_firstly(self.boxes, np.median(self.mean_height) / 3)
         # is it English
-        eng = 0
-        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
+        eng = is_english([b["text"] for b in bxs])
         # Merge vertically
         i = 0
         while i + 1 < len(bxs):
@@ -59,7 +53,7 @@ class Pdf(HuParser):
                 bxs.pop(i)
                 continue
             concatting_feats = [
-                b["text"].strip()[-1] in ",;:'\"，、‘“；：",
+                b["text"].strip()[-1] in ",;:'\"，、‘“；：-",
                 len(b["text"].strip())>1 and b["text"].strip()[-2] 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]
 
     # is it English
-    eng = 0
-    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
+    eng = is_english(sections)
     # Remove 'Contents' part
     i = 0
     while i < len(sections):
@@ -181,8 +168,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, callback=None):
         if pdf_parser:
             d["image"] = pdf_parser.crop(ck)
             ck = pdf_parser.remove_tag(ck)
-        d["content_ltks"] = huqie.qie(ck)
-        d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
+        tokenize(d, ck, eng)
         res.append(d)
     return res
 

rag/app/manual.py

@@ -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])

rag/app/paper.py

@@ -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])

rag/app/presentation.py

@@ -3,7 +3,7 @@ import re
 from io import BytesIO
 from pptx import Presentation
 
-from rag.app import callback__
+from rag.app import callback__, tokenize, is_english
 from rag.nlp import huqie
 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))
         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)
-
+        self.is_english = is_english(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"])
     res = []
     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["content_ltks"] = huqie.qie(txt)
-            d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
             d["image"] = img
+            tokenize(d, txt, ppt_parser.is_english)
             res.append(d)
         return res
     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["content_ltks"] = huqie.qie(txt)
-            d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
             d["image"] = img
+            tokenize(d, txt, pdf_parser.is_english)
             res.append(d)
         return res
     callback__(-1, "This kind of presentation document did not support yet!", callback)