KMS/ragflow
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

adjust hierarchical_merge strategy (#100)

Browse Source
This commit is contained in:
KevinHuSh 2024-03-06 09:09:16 +08:00 committed by GitHub
parent 8b2bbd79e4
commit d7c362f237
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 77 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
rag/app/laws.py
View File

@ -103,7 +103,6 @@ def chunk(filename, binary=None, from_page=0, to_page=100000, lang="Chinese", ca
if not l:break
txt += l
sections = txt.split("\n")
sections = txt.split("\n")
sections = [l for l in sections if l]
callback(0.8, "Finish parsing.")
else: raise NotImplementedError("file type not supported yet(docx, pdf, txt supported)")

117
rag/nlp/__init__.py
View File

@ -1,13 +1,14 @@
import random
from rag.utils import num_tokens_from_string
from . import huqie
from nltk import word_tokenize
import re
import copy
from nltk.stem import PorterStemmer
stemmer = PorterStemmer()
import re
from nltk import word_tokenize
from . import huqie
from rag.utils import num_tokens_from_string
import random
BULLET_PATTERN = [[
r"第[零一二三四五六七八九十百0-9]+(分?编|部分)",
@ -54,7 +55,8 @@ def bullets_category(sections):
maxium = 0
res = -1
for i, h in enumerate(hits):
if h <= maxium: continue
if h <= maxium:
continue
res = i
maxium = h
return res
@ -74,7 +76,8 @@ 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)])
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"])
@ -84,7 +87,8 @@ def tokenize_table(tbls, doc, eng, batch_size=10):
res = []
# add tables
for (img, rows), poss in tbls:
if not rows:continue
if not rows:
continue
if isinstance(rows, str):
d = copy.deepcopy(doc)
r = re.sub(r"<[^<>]{,12}>", "", rows)
@ -106,14 +110,15 @@ def tokenize_table(tbls, doc, eng, batch_size=10):
def add_positions(d, poss):
if not poss:return
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["page_num_int"].append(pn + 1)
d["top_int"].append(top)
d["position_int"].append((pn+1, left, right, top, bottom))
d["position_int"].append((pn + 1, left, right, top, bottom))
d["top_int"] = d["top_int"][:1]
@ -122,31 +127,38 @@ def remove_contents_table(sections, eng=False):
while i < len(sections):
def get(i):
nonlocal sections
return (sections[i] if type(sections[i]) == type("") else sections[i][0]).strip()
return (sections[i] if isinstance(sections[i],
type("")) else sections[i][0]).strip()
if not re.match(r"(contents|目录|目次|table of contents|致谢|acknowledge)$",
re.sub(r"( | |\u3000)+", "", get(i).split("@@")[0], re.IGNORECASE)):
i += 1
continue
sections.pop(i)
if i >= len(sections): break
if i >= len(sections):
break
prefix = get(i)[:3] if not eng else " ".join(get(i).split(" ")[:2])
while not prefix:
sections.pop(i)
if i >= len(sections): break
if i >= len(sections):
break
prefix = get(i)[:3] if not eng else " ".join(get(i).split(" ")[:2])
sections.pop(i)
if i >= len(sections) or not prefix: break
if i >= len(sections) or not prefix:
break
for j in range(i, min(i + 128, len(sections))):
if not re.match(prefix, get(j)):
continue
for _ in range(i, j): sections.pop(i)
for _ in range(i, j):
sections.pop(i)
break
def make_colon_as_title(sections):
if not sections: return []
if type(sections[0]) == type(""): return sections
if not sections:
return []
if isinstance(sections[0], type("")):
return sections
i = 0
while i < len(sections):
txt, layout = sections[i]
@ -165,20 +177,25 @@ def make_colon_as_title(sections):
def hierarchical_merge(bull, sections, depth):
if not sections or bull < 0: return []
if type(sections[0]) == type(""): sections = [(s, "") for s in sections]
sections = [(t,o) for t, o in sections if t and len(t.split("@")[0].strip()) > 1 and not re.match(r"[0-9]+$", t.split("@")[0].strip())]
if not sections or bull < 0:
return []
if isinstance(sections[0], type("")):
sections = [(s, "") for s in sections]
sections = [(t, o) for t, o in sections if
t and len(t.split("@")[0].strip()) > 1 and not re.match(r"[0-9]+$", t.split("@")[0].strip())]
bullets_size = len(BULLET_PATTERN[bull])
levels = [[] for _ in range(bullets_size + 2)]
def not_title(txt):
if re.match(r"第[零一二三四五六七八九十百0-9]+条", txt): return False
if len(txt.split(" "))>12 or (txt.find(" ")<0 and len(txt)) >= 32: return True
if re.match(r"第[零一二三四五六七八九十百0-9]+条", txt):
return False
if len(txt.split(" ")) > 12 or (txt.find(" ") < 0 and len(txt) >= 32):
return True
return re.search(r"[,;，。；！!]", txt)
for i, (txt, layout) in enumerate(sections):
for j, p in enumerate(BULLET_PATTERN[bull]):
if re.match(p, txt.strip()) and not not_title(txt):
if re.match(p, txt.strip()):
levels[j].append(i)
break
else:
@ -187,12 +204,16 @@ def hierarchical_merge(bull, sections, depth):
else:
levels[bullets_size + 1].append(i)
sections = [t for t, _ in sections]
#for s in sections: print("--", s)
# for s in sections: print("--", s)
def binary_search(arr, target):
if not arr: return -1
if target > arr[-1]: return len(arr) - 1
if target < arr[0]: return -1
if not arr:
return -1
if target > arr[-1]:
return len(arr) - 1
if target < arr[0]:
return -1
s, e = 0, len(arr)
while e - s > 1:
i = (e + s) // 2
@ -211,18 +232,24 @@ def hierarchical_merge(bull, sections, depth):
levels = levels[::-1]
for i, arr in enumerate(levels[:depth]):
for j in arr:
if readed[j]: continue
if readed[j]:
continue
readed[j] = True
cks.append([j])
if i + 1 == len(levels) - 1: continue
if i + 1 == len(levels) - 1:
continue
for ii in range(i + 1, len(levels)):
jj = binary_search(levels[ii], j)
if jj < 0: continue
if jj > cks[-1][-1]: cks[-1].pop(-1)
if jj < 0:
continue
if jj > cks[-1][-1]:
cks[-1].pop(-1)
cks[-1].append(levels[ii][jj])
for ii in cks[-1]: readed[ii] = True
for ii in cks[-1]:
readed[ii] = True
if not cks:return cks
if not cks:
return cks
for i in range(len(cks)):
cks[i] = [sections[j] for j in cks[i][::-1]]
@ -247,20 +274,26 @@ def hierarchical_merge(bull, sections, depth):
def naive_merge(sections, chunk_token_num=128, delimiter="\n。；！？"):
if not sections: return []
if type(sections[0]) == type(""): sections = [(s, "") for s in sections]
if not sections:
return []
if isinstance(sections[0], type("")):
sections = [(s, "") for s in sections]
cks = [""]
tk_nums = [0]
def add_chunk(t, pos):
nonlocal cks, tk_nums, delimiter
tnum = num_tokens_from_string(t)
if tnum < 8: pos = ""
if tnum < 8:
pos = ""
if tk_nums[-1] > chunk_token_num:
if t.find(pos) < 0: t += pos
if t.find(pos) < 0:
t += pos
cks.append(t)
tk_nums.append(tnum)
else:
if cks[-1].find(pos) < 0: t += pos
if cks[-1].find(pos) < 0:
t += pos
cks[-1] += t
tk_nums[-1] += tnum
@ -270,12 +303,12 @@ def naive_merge(sections, chunk_token_num=128, delimiter="\n。；！？"):
s, e = 0, 1
while e < len(sec):
if sec[e] in delimiter:
add_chunk(sec[s: e+1], pos)
add_chunk(sec[s: e + 1], pos)
s = e + 1
e = s + 1
else:
e += 1
if s < e: add_chunk(sec[s: e], pos)
if s < e:
add_chunk(sec[s: e], pos)
return cks

4
rag/nlp/search.py
View File

@ -82,8 +82,8 @@ class Dealer:
)
else:
s = s.sort(
{"page_num_int": {"order": "asc", "unmapped_type": "float"}},
{"top_int": {"order": "asc", "unmapped_type": "float", "mode" : "avg"}},
{"page_num_int": {"order": "asc", "unmapped_type": "float", "mode" : "avg"}},
{"top_int": {"order": "asc", "unmapped_type": "float", "mode": "avg"}},
{"create_time": {"order": "desc", "unmapped_type": "date"}},
{"create_timestamp_flt": {"order": "desc", "unmapped_type": "float"}}
)