ragflow/python/nlp/search.py

import re
from elasticsearch_dsl import Q,Search,A
from typing import List, Optional, Tuple,Dict, Union
from dataclasses import dataclass
from util import setup_logging, rmSpace
from nlp import huqie, query
from datetime import datetime
from sklearn.metrics.pairwise import cosine_similarity as CosineSimilarity
import numpy as np
from copy import deepcopy

def index_name(uid):return f"docgpt_{uid}"

class Dealer:
    def __init__(self, es, emb_mdl):
        self.qryr = query.EsQueryer(es)
        self.qryr.flds = ["title_tks^10", "title_sm_tks^5", "content_ltks^2", "content_sm_ltks"]
        self.es = es
        self.emb_mdl = emb_mdl

    @dataclass
    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, sim=0.8, topk=10):
        return {
            "field": "q_vec",
            "k": topk,
            "similarity": sim,
            "num_candidates": 1000,
            "query_vector": self.emb_mdl.encode_queries(txt)
        }

    def search(self, req, idxnm, tks_num=3):
        keywords = []
        qst = req.get("question", "")

        bqry,keywords = self.qryr.question(qst)
        if req.get("kb_ids"): bqry.filter.append(Q("terms", kb_id=req["kb_ids"]))
        bqry.filter.append(Q("exists", field="q_tks"))
        bqry.boost = 0.05
        print(bqry)

        s = Search()
        pg = int(req.get("page", 1))-1
        ps = int(req.get("size", 1000))
        src = req.get("field", ["docnm_kwd", "content_ltks", "kb_id",
                                "image_id", "doc_id", "q_vec"])

        s = s.query(bqry)[pg*ps:(pg+1)*ps]
        s = s.highlight("content_ltks")
        s = s.highlight("title_ltks")
        if not qst: s = s.sort({"create_time":{"order":"desc", "unmapped_type":"date"}})

        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"): 
            s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps)
            s["knn"]["filter"] = bqry.to_dict()
            del s["highlight"]
            q_vec = s["knn"]["query_vector"]
        res = self.es.search(s, idxnm=idxnm, timeout="600s",src=src)
        print("TOTAL: ", self.es.getTotal(res))
        if self.es.getTotal(res) == 0 and "knn" in s:
            bqry,_ = self.qryr.question(qst, min_match="10%")
            if req.get("kb_ids"): bqry.filter.append(Q("terms", kb_id=req["kb_ids"]))
            s["query"] = bqry.to_dict()
            s["knn"]["filter"] = bqry.to_dict()
            s["knn"]["similarity"] = 0.7
            res = self.es.search(s, idxnm=idxnm, timeout="600s",src=src)

        kwds = set([])
        for k in keywords:
            kwds.add(k)
            for kk in huqie.qieqie(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, ["docnm_kwd", "content_ltks",
                                         "kb_id","image_id", "doc_id", "q_vec"]),
            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 type(v) == type([]):
                    m[n] = "\t".join([str(vv) for vv in v])
                    continue
                if type(v) != type(""):m[n] = str(m[n])
                m[n] = rmSpace(m[n])

            if m:res[d["id"]] = m
        return res


    @staticmethod
    def trans2floats(txt):
        return [float(t) for t in txt.split("\t")]


    def insert_citations(self, ans, top_idx, sres, vfield = "q_vec", cfield="content_ltks"):

        ins_embd = [Dealer.trans2floats(sres.field[sres.ids[i]][vfield]) for i in top_idx]
        ins_tw =[sres.field[sres.ids[i]][cfield].split(" ") for i in top_idx]
        s = 0
        e = 0
        res = ""
        def citeit():
            nonlocal s, e, ans, res
            if not ins_embd:return
            embd = self.emb_mdl.encode(ans[s: e])
            sim = self.qryr.hybrid_similarity(embd, 
                                              ins_embd, 
                                              huqie.qie(ans[s:e]).split(" "),
                                              ins_tw)
            print(ans[s: e], sim)
            mx = np.max(sim)*0.99
            if mx < 0.55:return
            cita = list(set([top_idx[i] for i in range(len(ins_embd)) if sim[i] >mx]))[:4]
            for i in cita: res += f"@?{i}?@"

            return cita

        punct = set("；。？!！")
        if not self.qryr.isChinese(ans): 
            punct.add("?")
            punct.add(".")
        while e < len(ans):
            if e - s < 12 or ans[e] not in punct:
                e += 1
                continue
            if ans[e] == "." and e+1<len(ans) and re.match(r"[0-9]", ans[e+1]):
                e += 1
                continue
            if ans[e] == "." and e-2>=0 and  ans[e-2] == "\n":
                e += 1
                continue
            res += ans[s: e]
            citeit()
            res += ans[e]
            e += 1
            s = e

        if s< len(ans): 
            res += ans[s:]
            citeit()

        return res


    def rerank(self, sres, query, tkweight=0.3, vtweight=0.7, vfield="q_vec", cfield="content_ltks"):
        ins_embd = [Dealer.trans2floats(sres.field[i]["q_vec"]) for i in sres.ids]
        if not ins_embd: return []
        ins_tw =[sres.field[i][cfield].split(" ") for i in sres.ids]
        #return CosineSimilarity([sres.query_vector], ins_embd)[0]
        sim = self.qryr.hybrid_similarity(sres.query_vector, 
                                          ins_embd, 
                                          huqie.qie(query).split(" "),
                                          ins_tw, tkweight, vtweight)
        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_*"))