Kevin Hu
commited on
Commit
·
6d597a0
1
Parent(s):
acd1df1
debug backend API for TAB 'search' (#2389)
Browse files### What problem does this PR solve?
#2247
### Type of change
- [x] New Feature (non-breaking change which adds functionality)
- api/apps/chunk_app.py +1 -1
- api/apps/conversation_app.py +3 -1
- api/db/services/dialog_service.py +1 -2
- graphrag/mind_map_extractor.py +1 -1
- rag/llm/embedding_model.py +1 -1
- rag/nlp/search.py +23 -14
api/apps/chunk_app.py
CHANGED
@@ -261,7 +261,7 @@ def retrieval_test():
|
|
261 |
kb_id = req["kb_id"]
|
262 |
if isinstance(kb_id, str): kb_id = [kb_id]
|
263 |
doc_ids = req.get("doc_ids", [])
|
264 |
-
similarity_threshold = float(req.get("similarity_threshold", 0.
|
265 |
vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
|
266 |
top = int(req.get("top_k", 1024))
|
267 |
|
|
|
261 |
kb_id = req["kb_id"]
|
262 |
if isinstance(kb_id, str): kb_id = [kb_id]
|
263 |
doc_ids = req.get("doc_ids", [])
|
264 |
+
similarity_threshold = float(req.get("similarity_threshold", 0.0))
|
265 |
vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
|
266 |
top = int(req.get("top_k", 1024))
|
267 |
|
api/apps/conversation_app.py
CHANGED
@@ -15,8 +15,8 @@
|
|
15 |
#
|
16 |
import json
|
17 |
import re
|
|
|
18 |
from copy import deepcopy
|
19 |
-
|
20 |
from api.db.services.user_service import UserTenantService
|
21 |
from flask import request, Response
|
22 |
from flask_login import login_required, current_user
|
@@ -333,6 +333,8 @@ def mindmap():
|
|
333 |
0.3, 0.3, aggs=False)
|
334 |
mindmap = MindMapExtractor(chat_mdl)
|
335 |
mind_map = mindmap([c["content_with_weight"] for c in ranks["chunks"]]).output
|
|
|
|
|
336 |
return get_json_result(data=mind_map)
|
337 |
|
338 |
|
|
|
15 |
#
|
16 |
import json
|
17 |
import re
|
18 |
+
import traceback
|
19 |
from copy import deepcopy
|
|
|
20 |
from api.db.services.user_service import UserTenantService
|
21 |
from flask import request, Response
|
22 |
from flask_login import login_required, current_user
|
|
|
333 |
0.3, 0.3, aggs=False)
|
334 |
mindmap = MindMapExtractor(chat_mdl)
|
335 |
mind_map = mindmap([c["content_with_weight"] for c in ranks["chunks"]]).output
|
336 |
+
if "error" in mind_map:
|
337 |
+
return server_error_response(Exception(mind_map["error"]))
|
338 |
return get_json_result(data=mind_map)
|
339 |
|
340 |
|
api/db/services/dialog_service.py
CHANGED
@@ -218,7 +218,7 @@ def chat(dialog, messages, stream=True, **kwargs):
|
|
218 |
for ans in chat_mdl.chat_streamly(prompt, msg[1:], gen_conf):
|
219 |
answer = ans
|
220 |
delta_ans = ans[len(last_ans):]
|
221 |
-
if num_tokens_from_string(delta_ans) <
|
222 |
continue
|
223 |
last_ans = answer
|
224 |
yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
|
@@ -404,7 +404,6 @@ def rewrite(tenant_id, llm_id, question):
|
|
404 |
|
405 |
|
406 |
def tts(tts_mdl, text):
|
407 |
-
return
|
408 |
if not tts_mdl or not text: return
|
409 |
bin = b""
|
410 |
for chunk in tts_mdl.tts(text):
|
|
|
218 |
for ans in chat_mdl.chat_streamly(prompt, msg[1:], gen_conf):
|
219 |
answer = ans
|
220 |
delta_ans = ans[len(last_ans):]
|
221 |
+
if num_tokens_from_string(delta_ans) < 16:
|
222 |
continue
|
223 |
last_ans = answer
|
224 |
yield {"answer": answer, "reference": {}, "audio_binary": tts(tts_mdl, delta_ans)}
|
|
|
404 |
|
405 |
|
406 |
def tts(tts_mdl, text):
|
|
|
407 |
if not tts_mdl or not text: return
|
408 |
bin = b""
|
409 |
for chunk in tts_mdl.tts(text):
|
graphrag/mind_map_extractor.py
CHANGED
@@ -107,7 +107,7 @@ class MindMapExtractor:
|
|
107 |
res.append(_.result())
|
108 |
|
109 |
if not res:
|
110 |
-
return MindMapResult(output={"root":
|
111 |
|
112 |
merge_json = reduce(self._merge, res)
|
113 |
if len(merge_json.keys()) > 1:
|
|
|
107 |
res.append(_.result())
|
108 |
|
109 |
if not res:
|
110 |
+
return MindMapResult(output={"id": "root", "children": []})
|
111 |
|
112 |
merge_json = reduce(self._merge, res)
|
113 |
if len(merge_json.keys()) > 1:
|
rag/llm/embedding_model.py
CHANGED
@@ -15,7 +15,7 @@
|
|
15 |
#
|
16 |
import re
|
17 |
from typing import Optional
|
18 |
-
import
|
19 |
import requests
|
20 |
from huggingface_hub import snapshot_download
|
21 |
from openai.lib.azure import AzureOpenAI
|
|
|
15 |
#
|
16 |
import re
|
17 |
from typing import Optional
|
18 |
+
import threading
|
19 |
import requests
|
20 |
from huggingface_hub import snapshot_download
|
21 |
from openai.lib.azure import AzureOpenAI
|
rag/nlp/search.py
CHANGED
@@ -224,6 +224,8 @@ class Dealer:
|
|
224 |
def insert_citations(self, answer, chunks, chunk_v,
|
225 |
embd_mdl, tkweight=0.1, vtweight=0.9):
|
226 |
assert len(chunks) == len(chunk_v)
|
|
|
|
|
227 |
pieces = re.split(r"(```)", answer)
|
228 |
if len(pieces) >= 3:
|
229 |
i = 0
|
@@ -263,7 +265,7 @@ class Dealer:
|
|
263 |
|
264 |
ans_v, _ = embd_mdl.encode(pieces_)
|
265 |
assert len(ans_v[0]) == len(chunk_v[0]), "The dimension of query and chunk do not match: {} vs. {}".format(
|
266 |
-
|
267 |
|
268 |
chunks_tks = [rag_tokenizer.tokenize(self.qryr.rmWWW(ck)).split(" ")
|
269 |
for ck in chunks]
|
@@ -360,29 +362,33 @@ class Dealer:
|
|
360 |
ranks = {"total": 0, "chunks": [], "doc_aggs": {}}
|
361 |
if not question:
|
362 |
return ranks
|
363 |
-
|
|
|
364 |
"question": question, "vector": True, "topk": top,
|
365 |
"similarity": similarity_threshold,
|
366 |
"available_int": 1}
|
|
|
|
|
|
|
367 |
sres = self.search(req, index_name(tenant_id), embd_mdl, highlight)
|
|
|
368 |
|
369 |
-
if
|
370 |
-
|
371 |
-
|
|
|
|
|
|
|
|
|
|
|
372 |
else:
|
373 |
-
sim
|
374 |
-
|
375 |
-
idx = np.argsort(sim * -1)
|
376 |
|
377 |
dim = len(sres.query_vector)
|
378 |
-
start_idx = (page - 1) * page_size
|
379 |
for i in idx:
|
380 |
if sim[i] < similarity_threshold:
|
381 |
break
|
382 |
-
ranks["total"] += 1
|
383 |
-
start_idx -= 1
|
384 |
-
if start_idx >= 0:
|
385 |
-
continue
|
386 |
if len(ranks["chunks"]) >= page_size:
|
387 |
if aggs:
|
388 |
continue
|
@@ -406,7 +412,10 @@ class Dealer:
|
|
406 |
"positions": sres.field[id].get("position_int", "").split("\t")
|
407 |
}
|
408 |
if highlight:
|
409 |
-
|
|
|
|
|
|
|
410 |
if len(d["positions"]) % 5 == 0:
|
411 |
poss = []
|
412 |
for i in range(0, len(d["positions"]), 5):
|
|
|
224 |
def insert_citations(self, answer, chunks, chunk_v,
|
225 |
embd_mdl, tkweight=0.1, vtweight=0.9):
|
226 |
assert len(chunks) == len(chunk_v)
|
227 |
+
if not chunks:
|
228 |
+
return answer, set([])
|
229 |
pieces = re.split(r"(```)", answer)
|
230 |
if len(pieces) >= 3:
|
231 |
i = 0
|
|
|
265 |
|
266 |
ans_v, _ = embd_mdl.encode(pieces_)
|
267 |
assert len(ans_v[0]) == len(chunk_v[0]), "The dimension of query and chunk do not match: {} vs. {}".format(
|
268 |
+
len(ans_v[0]), len(chunk_v[0]))
|
269 |
|
270 |
chunks_tks = [rag_tokenizer.tokenize(self.qryr.rmWWW(ck)).split(" ")
|
271 |
for ck in chunks]
|
|
|
362 |
ranks = {"total": 0, "chunks": [], "doc_aggs": {}}
|
363 |
if not question:
|
364 |
return ranks
|
365 |
+
RERANK_PAGE_LIMIT = 3
|
366 |
+
req = {"kb_ids": kb_ids, "doc_ids": doc_ids, "size": page_size*RERANK_PAGE_LIMIT,
|
367 |
"question": question, "vector": True, "topk": top,
|
368 |
"similarity": similarity_threshold,
|
369 |
"available_int": 1}
|
370 |
+
if page > RERANK_PAGE_LIMIT:
|
371 |
+
req["page"] = page
|
372 |
+
req["size"] = page_size
|
373 |
sres = self.search(req, index_name(tenant_id), embd_mdl, highlight)
|
374 |
+
ranks["total"] = sres.total
|
375 |
|
376 |
+
if page <= RERANK_PAGE_LIMIT:
|
377 |
+
if rerank_mdl:
|
378 |
+
sim, tsim, vsim = self.rerank_by_model(rerank_mdl,
|
379 |
+
sres, question, 1 - vector_similarity_weight, vector_similarity_weight)
|
380 |
+
else:
|
381 |
+
sim, tsim, vsim = self.rerank(
|
382 |
+
sres, question, 1 - vector_similarity_weight, vector_similarity_weight)
|
383 |
+
idx = np.argsort(sim * -1)[(page-1)*page_size:page*page_size]
|
384 |
else:
|
385 |
+
sim = tsim = vsim = [1]*len(sres.ids)
|
386 |
+
idx = list(range(len(sres.ids)))
|
|
|
387 |
|
388 |
dim = len(sres.query_vector)
|
|
|
389 |
for i in idx:
|
390 |
if sim[i] < similarity_threshold:
|
391 |
break
|
|
|
|
|
|
|
|
|
392 |
if len(ranks["chunks"]) >= page_size:
|
393 |
if aggs:
|
394 |
continue
|
|
|
412 |
"positions": sres.field[id].get("position_int", "").split("\t")
|
413 |
}
|
414 |
if highlight:
|
415 |
+
if id in sres.highlight:
|
416 |
+
d["highlight"] = rmSpace(sres.highlight[id])
|
417 |
+
else:
|
418 |
+
d["highlight"] = d["content_with_weight"]
|
419 |
if len(d["positions"]) % 5 == 0:
|
420 |
poss = []
|
421 |
for i in range(0, len(d["positions"]), 5):
|