mirror of
https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git
synced 2025-04-20 05:00:01 +08:00
parent
6b8fc2ce1f
commit
c372afe40a
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
|
@ -1,8 +1,11 @@
|
|||||||
# -*- coding: utf-8 -*-
|
# -*- coding: utf-8 -*-
|
||||||
|
import json
|
||||||
import re
|
import re
|
||||||
from elasticsearch_dsl import Q, Search, A
|
from elasticsearch_dsl import Q, Search, A
|
||||||
from typing import List, Optional, Tuple, Dict, Union
|
from typing import List, Optional, Tuple, Dict, Union
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
|
from rag.settings import es_logger
|
||||||
from rag.utils import rmSpace
|
from rag.utils import rmSpace
|
||||||
from rag.nlp import huqie, query
|
from rag.nlp import huqie, query
|
||||||
import numpy as np
|
import numpy as np
|
||||||
@ -34,30 +37,30 @@ class Dealer:
|
|||||||
group_docs: List[List] = None
|
group_docs: List[List] = None
|
||||||
|
|
||||||
def _vector(self, txt, sim=0.8, topk=10):
|
def _vector(self, txt, sim=0.8, topk=10):
|
||||||
|
qv = self.emb_mdl.encode_queries(txt)
|
||||||
return {
|
return {
|
||||||
"field": "q_vec",
|
"field": "q_%d_vec"%len(qv),
|
||||||
"k": topk,
|
"k": topk,
|
||||||
"similarity": sim,
|
"similarity": sim,
|
||||||
"num_candidates": 1000,
|
"num_candidates": 1000,
|
||||||
"query_vector": self.emb_mdl.encode_queries(txt)
|
"query_vector": qv
|
||||||
}
|
}
|
||||||
|
|
||||||
def search(self, req, idxnm, tks_num=3):
|
def search(self, req, idxnm, tks_num=3):
|
||||||
keywords = []
|
|
||||||
qst = req.get("question", "")
|
qst = req.get("question", "")
|
||||||
|
|
||||||
bqry, keywords = self.qryr.question(qst)
|
bqry, keywords = self.qryr.question(qst)
|
||||||
if req.get("kb_ids"):
|
if req.get("kb_ids"):
|
||||||
bqry.filter.append(Q("terms", kb_id=req["kb_ids"]))
|
bqry.filter.append(Q("terms", kb_id=req["kb_ids"]))
|
||||||
bqry.filter.append(Q("exists", field="q_tks"))
|
if req.get("doc_ids"):
|
||||||
|
bqry.filter.append(Q("terms", doc_id=req["doc_ids"]))
|
||||||
bqry.boost = 0.05
|
bqry.boost = 0.05
|
||||||
print(bqry)
|
|
||||||
|
|
||||||
s = Search()
|
s = Search()
|
||||||
pg = int(req.get("page", 1)) - 1
|
pg = int(req.get("page", 1)) - 1
|
||||||
ps = int(req.get("size", 1000))
|
ps = int(req.get("size", 1000))
|
||||||
src = req.get("field", ["docnm_kwd", "content_ltks", "kb_id",
|
src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id","img_id",
|
||||||
"image_id", "doc_id", "q_vec"])
|
"image_id", "doc_id", "q_512_vec", "q_768_vec",
|
||||||
|
"q_1024_vec", "q_1536_vec"])
|
||||||
|
|
||||||
s = s.query(bqry)[pg * ps:(pg + 1) * ps]
|
s = s.query(bqry)[pg * ps:(pg + 1) * ps]
|
||||||
s = s.highlight("content_ltks")
|
s = s.highlight("content_ltks")
|
||||||
@ -66,22 +69,24 @@ class Dealer:
|
|||||||
s = s.sort(
|
s = s.sort(
|
||||||
{"create_time": {"order": "desc", "unmapped_type": "date"}})
|
{"create_time": {"order": "desc", "unmapped_type": "date"}})
|
||||||
|
|
||||||
s = s.highlight_options(
|
if qst:
|
||||||
fragment_size=120,
|
s = s.highlight_options(
|
||||||
number_of_fragments=5,
|
fragment_size=120,
|
||||||
boundary_scanner_locale="zh-CN",
|
number_of_fragments=5,
|
||||||
boundary_scanner="SENTENCE",
|
boundary_scanner_locale="zh-CN",
|
||||||
boundary_chars=",./;:\\!(),。?:!……()——、"
|
boundary_scanner="SENTENCE",
|
||||||
)
|
boundary_chars=",./;:\\!(),。?:!……()——、"
|
||||||
|
)
|
||||||
s = s.to_dict()
|
s = s.to_dict()
|
||||||
q_vec = []
|
q_vec = []
|
||||||
if req.get("vector"):
|
if req.get("vector"):
|
||||||
s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps)
|
s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps)
|
||||||
s["knn"]["filter"] = bqry.to_dict()
|
s["knn"]["filter"] = bqry.to_dict()
|
||||||
del s["highlight"]
|
if "highlight" in s: del s["highlight"]
|
||||||
q_vec = s["knn"]["query_vector"]
|
q_vec = s["knn"]["query_vector"]
|
||||||
|
es_logger.info("【Q】: {}".format(json.dumps(s)))
|
||||||
res = self.es.search(s, idxnm=idxnm, timeout="600s", src=src)
|
res = self.es.search(s, idxnm=idxnm, timeout="600s", src=src)
|
||||||
print("TOTAL: ", self.es.getTotal(res))
|
es_logger.info("TOTAL: {}".format(self.es.getTotal(res)))
|
||||||
if self.es.getTotal(res) == 0 and "knn" in s:
|
if self.es.getTotal(res) == 0 and "knn" in s:
|
||||||
bqry, _ = self.qryr.question(qst, min_match="10%")
|
bqry, _ = self.qryr.question(qst, min_match="10%")
|
||||||
if req.get("kb_ids"):
|
if req.get("kb_ids"):
|
||||||
@ -109,8 +114,7 @@ class Dealer:
|
|||||||
query_vector=q_vec,
|
query_vector=q_vec,
|
||||||
aggregation=aggs,
|
aggregation=aggs,
|
||||||
highlight=self.getHighlight(res),
|
highlight=self.getHighlight(res),
|
||||||
field=self.getFields(res, ["docnm_kwd", "content_ltks",
|
field=self.getFields(res, src),
|
||||||
"kb_id", "image_id", "doc_id", "q_vec"]),
|
|
||||||
keywords=list(kwds)
|
keywords=list(kwds)
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -237,14 +241,4 @@ class Dealer:
|
|||||||
return sim
|
return sim
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
from util import es_conn
|
|
||||||
SE = Dealer(es_conn.HuEs("infiniflow"))
|
|
||||||
qs = [
|
|
||||||
"胡凯",
|
|
||||||
""
|
|
||||||
]
|
|
||||||
for q in qs:
|
|
||||||
print(">>>>>>>>>>>>>>>>>>>>", q)
|
|
||||||
print(SE.search(
|
|
||||||
{"question": q, "kb_ids": "64f072a75f3b97c865718c4a"}, "infiniflow_*"))
|
|
||||||
|
@ -62,7 +62,7 @@ class Dealer:
|
|||||||
return set(res.keys())
|
return set(res.keys())
|
||||||
return res
|
return res
|
||||||
|
|
||||||
fnm = os.path.join(get_project_base_directory(), "res")
|
fnm = os.path.join(get_project_base_directory(), "rag/res")
|
||||||
self.ne, self.df = {}, {}
|
self.ne, self.df = {}, {}
|
||||||
try:
|
try:
|
||||||
self.ne = json.load(open(os.path.join(fnm, "ner.json"), "r"))
|
self.ne = json.load(open(os.path.join(fnm, "ner.json"), "r"))
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
|
@ -1,5 +1,5 @@
|
|||||||
#
|
#
|
||||||
# Copyright 2019 The FATE Authors. All Rights Reserved.
|
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
|
||||||
#
|
#
|
||||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
# you may not use this file except in compliance with the License.
|
# you may not use this file except in compliance with the License.
|
||||||
@ -13,6 +13,7 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
#
|
#
|
||||||
|
import datetime
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
@ -108,17 +109,17 @@ def build(row, cvmdl):
|
|||||||
(int(DOC_MAXIMUM_SIZE / 1024 / 1024)))
|
(int(DOC_MAXIMUM_SIZE / 1024 / 1024)))
|
||||||
return []
|
return []
|
||||||
|
|
||||||
res = ELASTICSEARCH.search(Q("term", doc_id=row["id"]))
|
# res = ELASTICSEARCH.search(Q("term", doc_id=row["id"]))
|
||||||
if ELASTICSEARCH.getTotal(res) > 0:
|
# if ELASTICSEARCH.getTotal(res) > 0:
|
||||||
ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]),
|
# ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]),
|
||||||
scripts="""
|
# scripts="""
|
||||||
if(!ctx._source.kb_id.contains('%s'))
|
# if(!ctx._source.kb_id.contains('%s'))
|
||||||
ctx._source.kb_id.add('%s');
|
# ctx._source.kb_id.add('%s');
|
||||||
""" % (str(row["kb_id"]), str(row["kb_id"])),
|
# """ % (str(row["kb_id"]), str(row["kb_id"])),
|
||||||
idxnm=search.index_name(row["tenant_id"])
|
# idxnm=search.index_name(row["tenant_id"])
|
||||||
)
|
# )
|
||||||
set_progress(row["id"], 1, "Done")
|
# set_progress(row["id"], 1, "Done")
|
||||||
return []
|
# return []
|
||||||
|
|
||||||
random.seed(time.time())
|
random.seed(time.time())
|
||||||
set_progress(row["id"], random.randint(0, 20) /
|
set_progress(row["id"], random.randint(0, 20) /
|
||||||
@ -155,8 +156,7 @@ def build(row, cvmdl):
|
|||||||
"doc_id": row["id"],
|
"doc_id": row["id"],
|
||||||
"kb_id": [str(row["kb_id"])],
|
"kb_id": [str(row["kb_id"])],
|
||||||
"docnm_kwd": os.path.split(row["location"])[-1],
|
"docnm_kwd": os.path.split(row["location"])[-1],
|
||||||
"title_tks": huqie.qie(row["name"]),
|
"title_tks": huqie.qie(row["name"])
|
||||||
"updated_at": str(row["update_time"]).replace("T", " ")[:19]
|
|
||||||
}
|
}
|
||||||
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"])
|
||||||
output_buffer = BytesIO()
|
output_buffer = BytesIO()
|
||||||
@ -179,6 +179,7 @@ def build(row, cvmdl):
|
|||||||
|
|
||||||
MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
|
MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
|
||||||
d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
|
d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
|
||||||
|
d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
|
||||||
docs.append(d)
|
docs.append(d)
|
||||||
|
|
||||||
for arr, img in obj.table_chunks:
|
for arr, img in obj.table_chunks:
|
||||||
@ -193,6 +194,7 @@ def build(row, cvmdl):
|
|||||||
img.save(output_buffer, format='JPEG')
|
img.save(output_buffer, format='JPEG')
|
||||||
MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
|
MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue())
|
||||||
d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
|
d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"])
|
||||||
|
d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19]
|
||||||
docs.append(d)
|
docs.append(d)
|
||||||
set_progress(row["id"], random.randint(60, 70) /
|
set_progress(row["id"], random.randint(60, 70) /
|
||||||
100., "Continue embedding the content.")
|
100., "Continue embedding the content.")
|
||||||
@ -218,23 +220,11 @@ def embedding(docs, mdl):
|
|||||||
vects = 0.1 * tts + 0.9 * cnts
|
vects = 0.1 * tts + 0.9 * cnts
|
||||||
assert len(vects) == len(docs)
|
assert len(vects) == len(docs)
|
||||||
for i, d in enumerate(docs):
|
for i, d in enumerate(docs):
|
||||||
d["q_vec"] = vects[i].tolist()
|
v = vects[i].tolist()
|
||||||
|
d["q_%d_vec"%len(v)] = v
|
||||||
return tk_count
|
return tk_count
|
||||||
|
|
||||||
|
|
||||||
def model_instance(tenant_id, llm_type):
|
|
||||||
model_config = TenantLLMService.get_api_key(tenant_id, model_type=LLMType.EMBEDDING)
|
|
||||||
if not model_config:
|
|
||||||
model_config = {"llm_factory": "local", "api_key": "", "llm_name": ""}
|
|
||||||
else: model_config = model_config[0].to_dict()
|
|
||||||
if llm_type == LLMType.EMBEDDING:
|
|
||||||
if model_config["llm_factory"] not in EmbeddingModel: return
|
|
||||||
return EmbeddingModel[model_config["llm_factory"]](model_config["api_key"], model_config["llm_name"])
|
|
||||||
if llm_type == LLMType.IMAGE2TEXT:
|
|
||||||
if model_config["llm_factory"] not in CvModel: return
|
|
||||||
return CvModel[model_config.llm_factory](model_config["api_key"], model_config["llm_name"])
|
|
||||||
|
|
||||||
|
|
||||||
def main(comm, mod):
|
def main(comm, mod):
|
||||||
global model
|
global model
|
||||||
from rag.llm import HuEmbedding
|
from rag.llm import HuEmbedding
|
||||||
@ -247,12 +237,12 @@ def main(comm, mod):
|
|||||||
|
|
||||||
tmf = open(tm_fnm, "a+")
|
tmf = open(tm_fnm, "a+")
|
||||||
for _, r in rows.iterrows():
|
for _, r in rows.iterrows():
|
||||||
embd_mdl = model_instance(r["tenant_id"], LLMType.EMBEDDING)
|
embd_mdl = TenantLLMService.model_instance(r["tenant_id"], LLMType.EMBEDDING)
|
||||||
if not embd_mdl:
|
if not embd_mdl:
|
||||||
set_progress(r["id"], -1, "Can't find embedding model!")
|
set_progress(r["id"], -1, "Can't find embedding model!")
|
||||||
cron_logger.error("Tenant({}) can't find embedding model!".format(r["tenant_id"]))
|
cron_logger.error("Tenant({}) can't find embedding model!".format(r["tenant_id"]))
|
||||||
continue
|
continue
|
||||||
cv_mdl = model_instance(r["tenant_id"], LLMType.IMAGE2TEXT)
|
cv_mdl = TenantLLMService.model_instance(r["tenant_id"], LLMType.IMAGE2TEXT)
|
||||||
st_tm = timer()
|
st_tm = timer()
|
||||||
cks = build(r, cv_mdl)
|
cks = build(r, cv_mdl)
|
||||||
if not cks:
|
if not cks:
|
||||||
|
@ -241,6 +241,26 @@ class HuEs:
|
|||||||
es_logger.error("ES search timeout for 3 times!")
|
es_logger.error("ES search timeout for 3 times!")
|
||||||
raise Exception("ES search timeout.")
|
raise Exception("ES search timeout.")
|
||||||
|
|
||||||
|
def get(self, doc_id, idxnm=None):
|
||||||
|
for i in range(3):
|
||||||
|
try:
|
||||||
|
res = self.es.get(index=(self.idxnm if not idxnm else idxnm),
|
||||||
|
id=doc_id)
|
||||||
|
if str(res.get("timed_out", "")).lower() == "true":
|
||||||
|
raise Exception("Es Timeout.")
|
||||||
|
return res
|
||||||
|
except Exception as e:
|
||||||
|
es_logger.error(
|
||||||
|
"ES get exception: " +
|
||||||
|
str(e) +
|
||||||
|
"【Q】:" +
|
||||||
|
doc_id)
|
||||||
|
if str(e).find("Timeout") > 0:
|
||||||
|
continue
|
||||||
|
raise e
|
||||||
|
es_logger.error("ES search timeout for 3 times!")
|
||||||
|
raise Exception("ES search timeout.")
|
||||||
|
|
||||||
def updateByQuery(self, q, d):
|
def updateByQuery(self, q, d):
|
||||||
ubq = UpdateByQuery(index=self.idxnm).using(self.es).query(q)
|
ubq = UpdateByQuery(index=self.idxnm).using(self.es).query(q)
|
||||||
scripts = ""
|
scripts = ""
|
||||||
|
Loading…
x
Reference in New Issue
Block a user