refine code (#595)

### What problem does this PR solve?

### Type of change

- [x] Refactoring

api/apps/chunk_app.py

@@ -20,7 +20,7 @@
 from elasticsearch_dsl import Q
 
 from rag.app.qa import rmPrefix, beAdoc
-from rag.nlp import search, huqie
+from rag.nlp import search, rag_tokenizer
 from rag.utils.es_conn import ELASTICSEARCH
 from rag.utils import rmSpace
 from api.db import LLMType, ParserType
@@ -125,10 +125,10 @@ def set():
     d = {
         "id": req["chunk_id"],
         "content_with_weight": req["content_with_weight"]}
-    d["content_ltks"] = huqie.qie(req["content_with_weight"])
-    d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
+    d["content_ltks"] = rag_tokenizer.tokenize(req["content_with_weight"])
+    d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
     d["important_kwd"] = req["important_kwd"]
-    d["important_tks"] = huqie.qie(" ".join(req["important_kwd"]))
+    d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_kwd"]))
     if "available_int" in req:
         d["available_int"] = req["available_int"]
 
@@ -152,7 +152,7 @@ def set():
                     retmsg="Q&A must be separated by TAB/ENTER key.")
             q, a = rmPrefix(arr[0]), rmPrefix[arr[1]]
             d = beAdoc(d, arr[0], arr[1], not any(
-                [huqie.is_chinese(t) for t in q + a]))
+                [rag_tokenizer.is_chinese(t) for t in q + a]))
 
         v, c = embd_mdl.encode([doc.name, req["content_with_weight"]])
         v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
@@ -202,11 +202,11 @@ def create():
     md5 = hashlib.md5()
     md5.update((req["content_with_weight"] + req["doc_id"]).encode("utf-8"))
     chunck_id = md5.hexdigest()
-    d = {"id": chunck_id, "content_ltks": huqie.qie(req["content_with_weight"]),
+    d = {"id": chunck_id, "content_ltks": rag_tokenizer.tokenize(req["content_with_weight"]),
          "content_with_weight": req["content_with_weight"]}
-    d["content_sm_ltks"] = huqie.qieqie(d["content_ltks"])
+    d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
     d["important_kwd"] = req.get("important_kwd", [])
-    d["important_tks"] = huqie.qie(" ".join(req.get("important_kwd", [])))
+    d["important_tks"] = rag_tokenizer.tokenize(" ".join(req.get("important_kwd", [])))
     d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
     d["create_timestamp_flt"] = datetime.datetime.now().timestamp()
 

api/db/services/task_service.py

@@ -78,14 +78,13 @@ def get_ongoing_doc_name(cls):
             docs = cls.model.select(*[Document.id, Document.kb_id, Document.location, File.parent_id]) \
                 .join(Document, on=(cls.model.doc_id == Document.id)) \
                 .join(File2Document, on=(File2Document.document_id == Document.id), join_type=JOIN.LEFT_OUTER) \
-                .join(File, on=(File2Document.file_id == File.id)) \
+                .join(File, on=(File2Document.file_id == File.id), join_type=JOIN.LEFT_OUTER) \
                 .where(
                     Document.status == StatusEnum.VALID.value,
                     Document.run == TaskStatus.RUNNING.value,
                     ~(Document.type == FileType.VIRTUAL.value),
-                    cls.model.progress >= 0,
                     cls.model.progress < 1,
-                    cls.model.create_time >= current_timestamp() - 180000
+                    cls.model.create_time >= current_timestamp() - 1000 * 600
                 )
             docs = list(docs.dicts())
             if not docs: return []

deepdoc/parser/docx_parser.py

@@ -3,7 +3,7 @@
 import re
 import pandas as pd
 from collections import Counter
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from io import BytesIO
 
 
@@ -35,14 +35,14 @@ def blockType(b):
             for p, n in patt:
                 if re.search(p, b):
                     return n
-            tks = [t for t in huqie.qie(b).split(" ") if len(t) > 1]
+            tks = [t for t in rag_tokenizer.tokenize(b).split(" ") if len(t) > 1]
             if len(tks) > 3:
                 if len(tks) < 12:
                     return "Tx"
                 else:
                     return "Lx"
 
-            if len(tks) == 1 and huqie.tag(tks[0]) == "nr":
+            if len(tks) == 1 and rag_tokenizer.tag(tks[0]) == "nr":
                 return "Nr"
 
             return "Ot"

deepdoc/parser/pdf_parser.py

@@ -16,7 +16,7 @@
 
 from api.utils.file_utils import get_project_base_directory
 from deepdoc.vision import OCR, Recognizer, LayoutRecognizer, TableStructureRecognizer
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from copy import deepcopy
 from huggingface_hub import snapshot_download
 
@@ -95,13 +95,13 @@ def _updown_concat_features(self, up, down):
         h = max(self.__height(up), self.__height(down))
         y_dis = self._y_dis(up, down)
         LEN = 6
-        tks_down = huqie.qie(down["text"][:LEN]).split(" ")
-        tks_up = huqie.qie(up["text"][-LEN:]).split(" ")
+        tks_down = rag_tokenizer.tokenize(down["text"][:LEN]).split(" ")
+        tks_up = rag_tokenizer.tokenize(up["text"][-LEN:]).split(" ")
         tks_all = up["text"][-LEN:].strip() \
                   + (" " if re.match(r"[a-zA-Z0-9]+",
                                      up["text"][-1] + down["text"][0]) else "") \
                   + down["text"][:LEN].strip()
-        tks_all = huqie.qie(tks_all).split(" ")
+        tks_all = rag_tokenizer.tokenize(tks_all).split(" ")
         fea = [
             up.get("R", -1) == down.get("R", -1),
             y_dis / h,
@@ -142,8 +142,8 @@ def _updown_concat_features(self, up, down):
             tks_down[-1] == tks_up[-1],
             max(down["in_row"], up["in_row"]),
             abs(down["in_row"] - up["in_row"]),
-            len(tks_down) == 1 and huqie.tag(tks_down[0]).find("n") >= 0,
-            len(tks_up) == 1 and huqie.tag(tks_up[0]).find("n") >= 0
+            len(tks_down) == 1 and rag_tokenizer.tag(tks_down[0]).find("n") >= 0,
+            len(tks_up) == 1 and rag_tokenizer.tag(tks_up[0]).find("n") >= 0
         ]
         return fea
 
@@ -599,7 +599,7 @@ def _merge_with_same_bullet(self):
 
             if b["text"].strip()[0] != b_["text"].strip()[0] \
                     or b["text"].strip()[0].lower() in set("qwertyuopasdfghjklzxcvbnm") \
-                    or huqie.is_chinese(b["text"].strip()[0]) \
+                    or rag_tokenizer.is_chinese(b["text"].strip()[0]) \
                     or b["top"] > b_["bottom"]:
                 i += 1
                 continue

deepdoc/parser/resume/entities/corporations.py

@@ -1,6 +1,6 @@
 import re,json,os
 import pandas as pd
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from . import regions
 current_file_path = os.path.dirname(os.path.abspath(__file__))
 GOODS = pd.read_csv(os.path.join(current_file_path, "res/corp_baike_len.csv"), sep="\t", header=0).fillna(0)
@@ -22,14 +22,14 @@ def baike(cid, default_v=0):
 def corpNorm(nm, add_region=True):
     global CORP_TKS
     if not nm or type(nm)!=type(""):return ""
-    nm = huqie.tradi2simp(huqie.strQ2B(nm)).lower()
+    nm = rag_tokenizer.tradi2simp(rag_tokenizer.strQ2B(nm)).lower()
     nm = re.sub(r"&", "&", nm)
     nm = re.sub(r"[\(\)（）\+'\"\t \*\\【】-]+", " ", nm)
     nm = re.sub(r"([—-]+.*| +co\..*|corp\..*| +inc\..*| +ltd.*)", "", nm, 10000, re.IGNORECASE)
     nm = re.sub(r"(计算机|技术|(技术|科技|网络)*有限公司|公司|有限|研发中心|中国|总部)$", "", nm, 10000, re.IGNORECASE)
     if not nm or (len(nm)<5 and not regions.isName(nm[0:2])):return nm
 
-    tks = huqie.qie(nm).split(" ")
+    tks = rag_tokenizer.tokenize(nm).split(" ")
     reg = [t for i,t in enumerate(tks) if regions.isName(t) and (t != "中国" or i > 0)]
     nm = ""
     for t in tks:

deepdoc/parser/resume/step_two.py

@@ -3,7 +3,7 @@
     traceback, signal
 import numpy as np
 from deepdoc.parser.resume.entities import degrees, schools, corporations
-from rag.nlp import huqie, surname
+from rag.nlp import rag_tokenizer, surname
 from xpinyin import Pinyin
 from contextlib import contextmanager
 
@@ -83,7 +83,7 @@ def forEdu(cv):
         if n.get("school_name") and isinstance(n["school_name"], str):
             sch.append(re.sub(r"(211|985|重点大学|[,&;；-])", "", n["school_name"]))
             e["sch_nm_kwd"] = sch[-1]
-        fea.append(huqie.qieqie(huqie.qie(n.get("school_name", ""))).split(" ")[-1])
+        fea.append(rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(n.get("school_name", ""))).split(" ")[-1])
 
         if n.get("discipline_name") and isinstance(n["discipline_name"], str):
             maj.append(n["discipline_name"])
@@ -166,10 +166,10 @@ def forEdu(cv):
             if "tag_kwd" not in cv: cv["tag_kwd"] = []
             if "好学历" not in cv["tag_kwd"]: cv["tag_kwd"].append("好学历")
 
-    if cv.get("major_kwd"): cv["major_tks"] = huqie.qie(" ".join(maj))
-    if cv.get("school_name_kwd"): cv["school_name_tks"] = huqie.qie(" ".join(sch))
-    if cv.get("first_school_name_kwd"): cv["first_school_name_tks"] = huqie.qie(" ".join(fsch))
-    if cv.get("first_major_kwd"): cv["first_major_tks"] = huqie.qie(" ".join(fmaj))
+    if cv.get("major_kwd"): cv["major_tks"] = rag_tokenizer.tokenize(" ".join(maj))
+    if cv.get("school_name_kwd"): cv["school_name_tks"] = rag_tokenizer.tokenize(" ".join(sch))
+    if cv.get("first_school_name_kwd"): cv["first_school_name_tks"] = rag_tokenizer.tokenize(" ".join(fsch))
+    if cv.get("first_major_kwd"): cv["first_major_tks"] = rag_tokenizer.tokenize(" ".join(fmaj))
 
     return cv
 
@@ -187,11 +187,11 @@ def forProj(cv):
         if n.get("achivement"): desc.append(str(n["achivement"]))
 
     if pro_nms:
-        # cv["pro_nms_tks"] = huqie.qie(" ".join(pro_nms))
-        cv["project_name_tks"] = huqie.qie(pro_nms[0])
+        # cv["pro_nms_tks"] = rag_tokenizer.tokenize(" ".join(pro_nms))
+        cv["project_name_tks"] = rag_tokenizer.tokenize(pro_nms[0])
     if desc:
-        cv["pro_desc_ltks"] = huqie.qie(rmHtmlTag(" ".join(desc)))
-        cv["project_desc_ltks"] = huqie.qie(rmHtmlTag(desc[0]))
+        cv["pro_desc_ltks"] = rag_tokenizer.tokenize(rmHtmlTag(" ".join(desc)))
+        cv["project_desc_ltks"] = rag_tokenizer.tokenize(rmHtmlTag(desc[0]))
 
     return cv
 
@@ -280,25 +280,25 @@ def forWork(cv):
     if fea["corporation_id"]: cv["corporation_id"] = fea["corporation_id"]
 
     if fea["position_name"]:
-        cv["position_name_tks"] = huqie.qie(fea["position_name"][0])
-        cv["position_name_sm_tks"] = huqie.qieqie(cv["position_name_tks"])
-        cv["pos_nm_tks"] = huqie.qie(" ".join(fea["position_name"][1:]))
+        cv["position_name_tks"] = rag_tokenizer.tokenize(fea["position_name"][0])
+        cv["position_name_sm_tks"] = rag_tokenizer.fine_grained_tokenize(cv["position_name_tks"])
+        cv["pos_nm_tks"] = rag_tokenizer.tokenize(" ".join(fea["position_name"][1:]))
 
     if fea["industry_name"]:
-        cv["industry_name_tks"] = huqie.qie(fea["industry_name"][0])
-        cv["industry_name_sm_tks"] = huqie.qieqie(cv["industry_name_tks"])
-        cv["indu_nm_tks"] = huqie.qie(" ".join(fea["industry_name"][1:]))
+        cv["industry_name_tks"] = rag_tokenizer.tokenize(fea["industry_name"][0])
+        cv["industry_name_sm_tks"] = rag_tokenizer.fine_grained_tokenize(cv["industry_name_tks"])
+        cv["indu_nm_tks"] = rag_tokenizer.tokenize(" ".join(fea["industry_name"][1:]))
 
     if fea["corporation_name"]:
         cv["corporation_name_kwd"] = fea["corporation_name"][0]
         cv["corp_nm_kwd"] = fea["corporation_name"]
-        cv["corporation_name_tks"] = huqie.qie(fea["corporation_name"][0])
-        cv["corporation_name_sm_tks"] = huqie.qieqie(cv["corporation_name_tks"])
-        cv["corp_nm_tks"] = huqie.qie(" ".join(fea["corporation_name"][1:]))
+        cv["corporation_name_tks"] = rag_tokenizer.tokenize(fea["corporation_name"][0])
+        cv["corporation_name_sm_tks"] = rag_tokenizer.fine_grained_tokenize(cv["corporation_name_tks"])
+        cv["corp_nm_tks"] = rag_tokenizer.tokenize(" ".join(fea["corporation_name"][1:]))
 
     if fea["responsibilities"]:
-        cv["responsibilities_ltks"] = huqie.qie(fea["responsibilities"][0])
-        cv["resp_ltks"] = huqie.qie(" ".join(fea["responsibilities"][1:]))
+        cv["responsibilities_ltks"] = rag_tokenizer.tokenize(fea["responsibilities"][0])
+        cv["resp_ltks"] = rag_tokenizer.tokenize(" ".join(fea["responsibilities"][1:]))
 
     if fea["subordinates_count"]: fea["subordinates_count"] = [int(i) for i in fea["subordinates_count"] if
                                                                re.match(r"[^0-9]+$", str(i))]
@@ -444,15 +444,15 @@ def hasValues(flds):
                 if nms:
                     t = k[:-4]
                     cv[f"{t}_kwd"] = nms
-                    cv[f"{t}_tks"] = huqie.qie(" ".join(nms))
+                    cv[f"{t}_tks"] = rag_tokenizer.tokenize(" ".join(nms))
             except Exception as e:
                 print("【EXCEPTION】:", str(traceback.format_exc()), cv[k])
                 cv[k] = []
 
         # tokenize fields
         if k in tks_fld:
-            cv[f"{k}_tks"] = huqie.qie(cv[k])
-            if k in small_tks_fld: cv[f"{k}_sm_tks"] = huqie.qie(cv[f"{k}_tks"])
+            cv[f"{k}_tks"] = rag_tokenizer.tokenize(cv[k])
+            if k in small_tks_fld: cv[f"{k}_sm_tks"] = rag_tokenizer.tokenize(cv[f"{k}_tks"])
 
         # keyword fields
         if k in kwd_fld: cv[f"{k}_kwd"] = [n.lower()
@@ -492,7 +492,7 @@ def hasValues(flds):
         cv["name_kwd"] = name
         cv["name_pinyin_kwd"] = PY.get_pinyins(nm[:20], ' ')[:3]
         cv["name_tks"] = (
-                huqie.qie(name) + " " + (" ".join(list(name)) if not re.match(r"[a-zA-Z ]+$", name) else "")
+                rag_tokenizer.tokenize(name) + " " + (" ".join(list(name)) if not re.match(r"[a-zA-Z ]+$", name) else "")
         ) if name else ""
     else:
         cv["integerity_flt"] /= 2.
@@ -515,7 +515,7 @@ def hasValues(flds):
         cv["updated_at_dt"] = f"%s-%02d-%02d 00:00:00" % (y, int(m), int(d))
         # long text tokenize
 
-    if cv.get("responsibilities"): cv["responsibilities_ltks"] = huqie.qie(rmHtmlTag(cv["responsibilities"]))
+    if cv.get("responsibilities"): cv["responsibilities_ltks"] = rag_tokenizer.tokenize(rmHtmlTag(cv["responsibilities"]))
 
     # for yes or no field
     fea = []

deepdoc/vision/table_structure_recognizer.py

@@ -19,7 +19,7 @@
 from huggingface_hub import snapshot_download
 
 from api.utils.file_utils import get_project_base_directory
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from .recognizer import Recognizer
 
 
@@ -117,14 +117,14 @@ def blockType(b):
         for p, n in patt:
             if re.search(p, b["text"].strip()):
                 return n
-        tks = [t for t in huqie.qie(b["text"]).split(" ") if len(t) > 1]
+        tks = [t for t in rag_tokenizer.tokenize(b["text"]).split(" ") if len(t) > 1]
         if len(tks) > 3:
             if len(tks) < 12:
                 return "Tx"
             else:
                 return "Lx"
 
-        if len(tks) == 1 and huqie.tag(tks[0]) == "nr":
+        if len(tks) == 1 and rag_tokenizer.tag(tks[0]) == "nr":
             return "Nr"
 
         return "Ot"

rag/app/book.py

@@ -18,7 +18,7 @@
 from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, \
     hierarchical_merge, make_colon_as_title, naive_merge, random_choices, tokenize_table, add_positions, \
     tokenize_chunks, find_codec
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from deepdoc.parser import PdfParser, DocxParser, PlainParser
 
 
@@ -63,9 +63,9 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
     """
     doc = {
         "docnm_kwd": filename,
-        "title_tks": huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", filename))
+        "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", filename))
     }
-    doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
+    doc["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(doc["title_tks"])
     pdf_parser = None
     sections, tbls = [], []
     if re.search(r"\.docx$", filename, re.IGNORECASE):

rag/app/laws.py

@@ -19,7 +19,7 @@
 from api.db import ParserType
 from rag.nlp import bullets_category, is_english, tokenize, remove_contents_table, hierarchical_merge, \
     make_colon_as_title, add_positions, tokenize_chunks, find_codec
-from rag.nlp import huqie
+from rag.nlp import rag_tokenizer
 from deepdoc.parser import PdfParser, DocxParser, PlainParser
 from rag.settings import cron_logger
 
@@ -89,9 +89,9 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
     """
     doc = {
         "docnm_kwd": filename,
-        "title_tks": huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", filename))
+        "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", filename))
     }
-    doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
+    doc["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(doc["title_tks"])
     pdf_parser = None
     sections = []
     if re.search(r"\.docx$", filename, re.IGNORECASE):

rag/app/manual.py

@@ -2,7 +2,7 @@
 import re
 
 from api.db import ParserType
-from rag.nlp import huqie, tokenize, tokenize_table, add_positions, bullets_category, title_frequency, tokenize_chunks
+from rag.nlp import rag_tokenizer, 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
 
@@ -70,8 +70,8 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
     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"])
+    doc["title_tks"] = rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", doc["docnm_kwd"]))
+    doc["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(doc["title_tks"])
     # is it English
     eng = lang.lower() == "english"  # pdf_parser.is_english
 

rag/app/naive.py

@@ -16,7 +16,7 @@
 from timeit import default_timer as timer
 import re
 from deepdoc.parser.pdf_parser import PlainParser
-from rag.nlp import huqie, naive_merge, tokenize_table, tokenize_chunks, find_codec
+from rag.nlp import rag_tokenizer, naive_merge, tokenize_table, tokenize_chunks, find_codec
 from deepdoc.parser import PdfParser, ExcelParser, DocxParser
 from rag.settings import cron_logger
 
@@ -112,9 +112,9 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
             "chunk_token_num": 128, "delimiter": "\n!?。；！？", "layout_recognize": True})
     doc = {
         "docnm_kwd": filename,
-        "title_tks": huqie.qie(re.sub(r"\.[a-zA-Z]+$", "", filename))
+        "title_tks": rag_tokenizer.tokenize(re.sub(r"\.[a-zA-Z]+$", "", filename))
     }
-    doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
+    doc["title_sm_tks"] = rag_tokenizer.fine_grained_tokenize(doc["title_tks"])
     res = []
     pdf_parser = None
     sections = []