Spaces:
Paused
Paused
| # | |
| # Copyright 2024 The InfiniFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # | |
| import json | |
| import re | |
| from copy import deepcopy | |
| from elasticsearch_dsl import Q, Search | |
| from typing import List, Optional, Dict, Union | |
| from dataclasses import dataclass | |
| from rag.settings import es_logger | |
| from rag.utils import rmSpace | |
| from rag.nlp import rag_tokenizer, query | |
| import numpy as np | |
| def index_name(uid): return f"ragflow_{uid}" | |
| class Dealer: | |
| def __init__(self, es): | |
| self.qryr = query.EsQueryer(es) | |
| self.qryr.flds = [ | |
| "title_tks^10", | |
| "title_sm_tks^5", | |
| "important_kwd^30", | |
| "important_tks^20", | |
| "content_ltks^2", | |
| "content_sm_ltks"] | |
| self.es = es | |
| class SearchResult: | |
| total: int | |
| ids: List[str] | |
| query_vector: List[float] = None | |
| field: Optional[Dict] = None | |
| highlight: Optional[Dict] = None | |
| aggregation: Union[List, Dict, None] = None | |
| keywords: Optional[List[str]] = None | |
| group_docs: List[List] = None | |
| def _vector(self, txt, emb_mdl, sim=0.8, topk=10): | |
| qv, c = emb_mdl.encode_queries(txt) | |
| return { | |
| "field": "q_%d_vec" % len(qv), | |
| "k": topk, | |
| "similarity": sim, | |
| "num_candidates": topk * 2, | |
| "query_vector": [float(v) for v in qv] | |
| } | |
| def _add_filters(self, bqry, req): | |
| if req.get("kb_ids"): | |
| bqry.filter.append(Q("terms", kb_id=req["kb_ids"])) | |
| if req.get("doc_ids"): | |
| bqry.filter.append(Q("terms", doc_id=req["doc_ids"])) | |
| if req.get("knowledge_graph_kwd"): | |
| bqry.filter.append(Q("terms", knowledge_graph_kwd=req["knowledge_graph_kwd"])) | |
| if "available_int" in req: | |
| if req["available_int"] == 0: | |
| bqry.filter.append(Q("range", available_int={"lt": 1})) | |
| else: | |
| bqry.filter.append( | |
| Q("bool", must_not=Q("range", available_int={"lt": 1}))) | |
| return bqry | |
| def search(self, req, idxnm, emb_mdl=None): | |
| qst = req.get("question", "") | |
| bqry, keywords = self.qryr.question(qst) | |
| bqry = self._add_filters(bqry, req) | |
| bqry.boost = 0.05 | |
| s = Search() | |
| pg = int(req.get("page", 1)) - 1 | |
| topk = int(req.get("topk", 1024)) | |
| ps = int(req.get("size", topk)) | |
| src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id", "img_id", "title_tks", "important_kwd", | |
| "image_id", "doc_id", "q_512_vec", "q_768_vec", "position_int", "knowledge_graph_kwd", | |
| "q_1024_vec", "q_1536_vec", "available_int", "content_with_weight"]) | |
| s = s.query(bqry)[pg * ps:(pg + 1) * ps] | |
| s = s.highlight("content_ltks") | |
| s = s.highlight("title_ltks") | |
| if not qst: | |
| if not req.get("sort"): | |
| s = s.sort( | |
| #{"create_time": {"order": "desc", "unmapped_type": "date"}}, | |
| {"create_timestamp_flt": { | |
| "order": "desc", "unmapped_type": "float"}} | |
| ) | |
| else: | |
| s = s.sort( | |
| {"page_num_int": {"order": "asc", "unmapped_type": "float", | |
| "mode": "avg", "numeric_type": "double"}}, | |
| {"top_int": {"order": "asc", "unmapped_type": "float", | |
| "mode": "avg", "numeric_type": "double"}}, | |
| #{"create_time": {"order": "desc", "unmapped_type": "date"}}, | |
| {"create_timestamp_flt": { | |
| "order": "desc", "unmapped_type": "float"}} | |
| ) | |
| if qst: | |
| s = s.highlight_options( | |
| fragment_size=120, | |
| number_of_fragments=5, | |
| boundary_scanner_locale="zh-CN", | |
| boundary_scanner="SENTENCE", | |
| boundary_chars=",./;:\\!()，。？：！……（）——、" | |
| ) | |
| s = s.to_dict() | |
| q_vec = [] | |
| if req.get("vector"): | |
| assert emb_mdl, "No embedding model selected" | |
| s["knn"] = self._vector( | |
| qst, emb_mdl, req.get( | |
| "similarity", 0.1), topk) | |
| s["knn"]["filter"] = bqry.to_dict() | |
| if "highlight" in s: | |
| del s["highlight"] | |
| q_vec = s["knn"]["query_vector"] | |
| es_logger.info("【Q】: {}".format(json.dumps(s))) | |
| res = self.es.search(deepcopy(s), idxnm=idxnm, timeout="600s", src=src) | |
| es_logger.info("TOTAL: {}".format(self.es.getTotal(res))) | |
| if self.es.getTotal(res) == 0 and "knn" in s: | |
| bqry, _ = self.qryr.question(qst, min_match="10%") | |
| bqry = self._add_filters(bqry, req) | |
| s["query"] = bqry.to_dict() | |
| s["knn"]["filter"] = bqry.to_dict() | |
| s["knn"]["similarity"] = 0.17 | |
| res = self.es.search(s, idxnm=idxnm, timeout="600s", src=src) | |
| es_logger.info("【Q】: {}".format(json.dumps(s))) | |
| kwds = set([]) | |
| for k in keywords: | |
| kwds.add(k) | |
| for kk in rag_tokenizer.fine_grained_tokenize(k).split(" "): | |
| if len(kk) < 2: | |
| continue | |
| if kk in kwds: | |
| continue | |
| kwds.add(kk) | |
| aggs = self.getAggregation(res, "docnm_kwd") | |
| return self.SearchResult( | |
| total=self.es.getTotal(res), | |
| ids=self.es.getDocIds(res), | |
| query_vector=q_vec, | |
| aggregation=aggs, | |
| highlight=self.getHighlight(res), | |
| field=self.getFields(res, src), | |
| keywords=list(kwds) | |
| ) | |
| def getAggregation(self, res, g): | |
| if not "aggregations" in res or "aggs_" + g not in res["aggregations"]: | |
| return | |
| bkts = res["aggregations"]["aggs_" + g]["buckets"] | |
| return [(b["key"], b["doc_count"]) for b in bkts] | |
| def getHighlight(self, res): | |
| def rmspace(line): | |
| eng = set(list("qwertyuioplkjhgfdsazxcvbnm")) | |
| r = [] | |
| for t in line.split(" "): | |
| if not t: | |
| continue | |
| if len(r) > 0 and len( | |
| t) > 0 and r[-1][-1] in eng and t[0] in eng: | |
| r.append(" ") | |
| r.append(t) | |
| r = "".join(r) | |
| return r | |
| ans = {} | |
| for d in res["hits"]["hits"]: | |
| hlts = d.get("highlight") | |
| if not hlts: | |
| continue | |
| ans[d["_id"]] = "".join([a for a in list(hlts.items())[0][1]]) | |
| return ans | |
| def getFields(self, sres, flds): | |
| res = {} | |
| if not flds: | |
| return {} | |
| for d in self.es.getSource(sres): | |
| m = {n: d.get(n) for n in flds if d.get(n) is not None} | |
| for n, v in m.items(): | |
| if isinstance(v, type([])): | |
| m[n] = "\t".join([str(vv) if not isinstance( | |
| vv, list) else "\t".join([str(vvv) for vvv in vv]) for vv in v]) | |
| continue | |
| if not isinstance(v, type("")): | |
| m[n] = str(m[n]) | |
| if n.find("tks") > 0: | |
| m[n] = rmSpace(m[n]) | |
| if m: | |
| res[d["id"]] = m | |
| return res | |
| def trans2floats(txt): | |
| return [float(t) for t in txt.split("\t")] | |
| def insert_citations(self, answer, chunks, chunk_v, | |
| embd_mdl, tkweight=0.1, vtweight=0.9): | |
| assert len(chunks) == len(chunk_v) | |
| pieces = re.split(r"(```)", answer) | |
| if len(pieces) >= 3: | |
| i = 0 | |
| pieces_ = [] | |
| while i < len(pieces): | |
| if pieces[i] == "```": | |
| st = i | |
| i += 1 | |
| while i < len(pieces) and pieces[i] != "```": | |
| i += 1 | |
| if i < len(pieces): | |
| i += 1 | |
| pieces_.append("".join(pieces[st: i]) + "\n") | |
| else: | |
| pieces_.extend( | |
| re.split( | |
| r"([^\|][；。？!！\n]|[a-z][.?;!][ \n])", | |
| pieces[i])) | |
| i += 1 | |
| pieces = pieces_ | |
| else: | |
| pieces = re.split(r"([^\|][；。？!！\n]|[a-z][.?;!][ \n])", answer) | |
| for i in range(1, len(pieces)): | |
| if re.match(r"([^\|][；。？!！\n]|[a-z][.?;!][ \n])", pieces[i]): | |
| pieces[i - 1] += pieces[i][0] | |
| pieces[i] = pieces[i][1:] | |
| idx = [] | |
| pieces_ = [] | |
| for i, t in enumerate(pieces): | |
| if len(t) < 5: | |
| continue | |
| idx.append(i) | |
| pieces_.append(t) | |
| es_logger.info("{} => {}".format(answer, pieces_)) | |
| if not pieces_: | |
| return answer, set([]) | |
| ans_v, _ = embd_mdl.encode(pieces_) | |
| assert len(ans_v[0]) == len(chunk_v[0]), "The dimension of query and chunk do not match: {} vs. {}".format( | |
| len(ans_v[0]), len(chunk_v[0])) | |
| chunks_tks = [rag_tokenizer.tokenize(self.qryr.rmWWW(ck)).split(" ") | |
| for ck in chunks] | |
| cites = {} | |
| thr = 0.63 | |
| while thr>0.3 and len(cites.keys()) == 0 and pieces_ and chunks_tks: | |
| for i, a in enumerate(pieces_): | |
| sim, tksim, vtsim = self.qryr.hybrid_similarity(ans_v[i], | |
| chunk_v, | |
| rag_tokenizer.tokenize( | |
| self.qryr.rmWWW(pieces_[i])).split(" "), | |
| chunks_tks, | |
| tkweight, vtweight) | |
| mx = np.max(sim) * 0.99 | |
| es_logger.info("{} SIM: {}".format(pieces_[i], mx)) | |
| if mx < thr: | |
| continue | |
| cites[idx[i]] = list( | |
| set([str(ii) for ii in range(len(chunk_v)) if sim[ii] > mx]))[:4] | |
| thr *= 0.8 | |
| res = "" | |
| seted = set([]) | |
| for i, p in enumerate(pieces): | |
| res += p | |
| if i not in idx: | |
| continue | |
| if i not in cites: | |
| continue | |
| for c in cites[i]: | |
| assert int(c) < len(chunk_v) | |
| for c in cites[i]: | |
| if c in seted: | |
| continue | |
| res += f" ##{c}$$" | |
| seted.add(c) | |
| return res, seted | |
| def rerank(self, sres, query, tkweight=0.3, | |
| vtweight=0.7, cfield="content_ltks"): | |
| _, keywords = self.qryr.question(query) | |
| ins_embd = [ | |
| Dealer.trans2floats( | |
| sres.field[i].get("q_%d_vec" % len(sres.query_vector), "\t".join(["0"] * len(sres.query_vector)))) for i in sres.ids] | |
| if not ins_embd: | |
| return [], [], [] | |
| for i in sres.ids: | |
| if isinstance(sres.field[i].get("important_kwd", []), str): | |
| sres.field[i]["important_kwd"] = [sres.field[i]["important_kwd"]] | |
| ins_tw = [] | |
| for i in sres.ids: | |
| content_ltks = sres.field[i][cfield].split(" ") | |
| title_tks = [t for t in sres.field[i].get("title_tks", "").split(" ") if t] | |
| important_kwd = sres.field[i].get("important_kwd", []) | |
| tks = content_ltks + title_tks + important_kwd | |
| ins_tw.append(tks) | |
| sim, tksim, vtsim = self.qryr.hybrid_similarity(sres.query_vector, | |
| ins_embd, | |
| keywords, | |
| ins_tw, tkweight, vtweight) | |
| return sim, tksim, vtsim | |
| def rerank_by_model(self, rerank_mdl, sres, query, tkweight=0.3, | |
| vtweight=0.7, cfield="content_ltks"): | |
| _, keywords = self.qryr.question(query) | |
| for i in sres.ids: | |
| if isinstance(sres.field[i].get("important_kwd", []), str): | |
| sres.field[i]["important_kwd"] = [sres.field[i]["important_kwd"]] | |
| ins_tw = [] | |
| for i in sres.ids: | |
| content_ltks = sres.field[i][cfield].split(" ") | |
| title_tks = [t for t in sres.field[i].get("title_tks", "").split(" ") if t] | |
| important_kwd = sres.field[i].get("important_kwd", []) | |
| tks = content_ltks + title_tks + important_kwd | |
| ins_tw.append(tks) | |
| tksim = self.qryr.token_similarity(keywords, ins_tw) | |
| vtsim,_ = rerank_mdl.similarity(" ".join(keywords), [rmSpace(" ".join(tks)) for tks in ins_tw]) | |
| return tkweight*np.array(tksim) + vtweight*vtsim, tksim, vtsim | |
| def hybrid_similarity(self, ans_embd, ins_embd, ans, inst): | |
| return self.qryr.hybrid_similarity(ans_embd, | |
| ins_embd, | |
| rag_tokenizer.tokenize(ans).split(" "), | |
| rag_tokenizer.tokenize(inst).split(" ")) | |
| def retrieval(self, question, embd_mdl, tenant_id, kb_ids, page, page_size, similarity_threshold=0.2, | |
| vector_similarity_weight=0.3, top=1024, doc_ids=None, aggs=True, rerank_mdl=None): | |
| ranks = {"total": 0, "chunks": [], "doc_aggs": {}} | |
| if not question: | |
| return ranks | |
| req = {"kb_ids": kb_ids, "doc_ids": doc_ids, "size": page_size, | |
| "question": question, "vector": True, "topk": top, | |
| "similarity": similarity_threshold, | |
| "available_int": 1} | |
| sres = self.search(req, index_name(tenant_id), embd_mdl) | |
| if rerank_mdl: | |
| sim, tsim, vsim = self.rerank_by_model(rerank_mdl, | |
| sres, question, 1 - vector_similarity_weight, vector_similarity_weight) | |
| else: | |
| sim, tsim, vsim = self.rerank( | |
| sres, question, 1 - vector_similarity_weight, vector_similarity_weight) | |
| idx = np.argsort(sim * -1) | |
| dim = len(sres.query_vector) | |
| start_idx = (page - 1) * page_size | |
| for i in idx: | |
| if sim[i] < similarity_threshold: | |
| break | |
| ranks["total"] += 1 | |
| start_idx -= 1 | |
| if start_idx >= 0: | |
| continue | |
| if len(ranks["chunks"]) >= page_size: | |
| if aggs: | |
| continue | |
| break | |
| id = sres.ids[i] | |
| dnm = sres.field[id]["docnm_kwd"] | |
| did = sres.field[id]["doc_id"] | |
| d = { | |
| "chunk_id": id, | |
| "content_ltks": sres.field[id]["content_ltks"], | |
| "content_with_weight": sres.field[id]["content_with_weight"], | |
| "doc_id": sres.field[id]["doc_id"], | |
| "docnm_kwd": dnm, | |
| "kb_id": sres.field[id]["kb_id"], | |
| "important_kwd": sres.field[id].get("important_kwd", []), | |
| "img_id": sres.field[id].get("img_id", ""), | |
| "similarity": sim[i], | |
| "vector_similarity": vsim[i], | |
| "term_similarity": tsim[i], | |
| "vector": self.trans2floats(sres.field[id].get("q_%d_vec" % dim, "\t".join(["0"] * dim))), | |
| "positions": sres.field[id].get("position_int", "").split("\t") | |
| } | |
| if len(d["positions"]) % 5 == 0: | |
| poss = [] | |
| for i in range(0, len(d["positions"]), 5): | |
| poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]), | |
| float(d["positions"][i + 3]), float(d["positions"][i + 4])]) | |
| d["positions"] = poss | |
| ranks["chunks"].append(d) | |
| if dnm not in ranks["doc_aggs"]: | |
| ranks["doc_aggs"][dnm] = {"doc_id": did, "count": 0} | |
| ranks["doc_aggs"][dnm]["count"] += 1 | |
| ranks["doc_aggs"] = [{"doc_name": k, | |
| "doc_id": v["doc_id"], | |
| "count": v["count"]} for k, | |
| v in sorted(ranks["doc_aggs"].items(), | |
| key=lambda x:x[1]["count"] * -1)] | |
| return ranks | |
| def sql_retrieval(self, sql, fetch_size=128, format="json"): | |
| from api.settings import chat_logger | |
| sql = re.sub(r"[ `]+", " ", sql) | |
| sql = sql.replace("%", "") | |
| es_logger.info(f"Get es sql: {sql}") | |
| replaces = [] | |
| for r in re.finditer(r" ([a-z_]+_l?tks)( like | ?= ?)'([^']+)'", sql): | |
| fld, v = r.group(1), r.group(3) | |
| match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format( | |
| fld, rag_tokenizer.fine_grained_tokenize(rag_tokenizer.tokenize(v))) | |
| replaces.append( | |
| ("{}{}'{}'".format( | |
| r.group(1), | |
| r.group(2), | |
| r.group(3)), | |
| match)) | |
| for p, r in replaces: | |
| sql = sql.replace(p, r, 1) | |
| chat_logger.info(f"To es: {sql}") | |
| try: | |
| tbl = self.es.sql(sql, fetch_size, format) | |
| return tbl | |
| except Exception as e: | |
| chat_logger.error(f"SQL failure: {sql} =>" + str(e)) | |
| return {"error": str(e)} | |
| def chunk_list(self, doc_id, tenant_id, max_count=1024, fields=["docnm_kwd", "content_with_weight", "img_id"]): | |
| s = Search() | |
| s = s.query(Q("match", doc_id=doc_id))[0:max_count] | |
| s = s.to_dict() | |
| es_res = self.es.search(s, idxnm=index_name(tenant_id), timeout="600s", src=fields) | |
| res = [] | |
| for index, chunk in enumerate(es_res['hits']['hits']): | |
| res.append({fld: chunk['_source'].get(fld) for fld in fields}) | |
| return res | |