From c372afe40a9c17172323474c06573be3cbeffa34 Mon Sep 17 00:00:00 2001 From: KevinHuSh Date: Wed, 17 Jan 2024 09:39:50 +0800 Subject: [PATCH] change licence (#28) * add front end code * change licence --- rag/llm/__init__.py | 2 +- rag/llm/chat_model.py | 2 +- rag/llm/cv_model.py | 2 +- rag/llm/embedding_model.py | 2 +- rag/nlp/search.py | 54 +++++++++++++++++--------------------- rag/nlp/term_weight.py | 2 +- rag/settings.py | 2 +- rag/svr/parse_user_docs.py | 50 ++++++++++++++--------------------- rag/utils/es_conn.py | 20 ++++++++++++++ 9 files changed, 70 insertions(+), 66 deletions(-) diff --git a/rag/llm/__init__.py b/rag/llm/__init__.py index 94a7ebe7c..b14a9db62 100644 --- a/rag/llm/__init__.py +++ b/rag/llm/__init__.py @@ -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"); # you may not use this file except in compliance with the License. diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 031aaab42..5564b9c49 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -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"); # you may not use this file except in compliance with the License. diff --git a/rag/llm/cv_model.py b/rag/llm/cv_model.py index 306cdf0c5..b3ec20290 100644 --- a/rag/llm/cv_model.py +++ b/rag/llm/cv_model.py @@ -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"); # you may not use this file except in compliance with the License. diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index e148560c2..20608720b 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -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"); # you may not use this file except in compliance with the License. diff --git a/rag/nlp/search.py b/rag/nlp/search.py index d79640b4c..442812c56 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -1,8 +1,11 @@ # -*- coding: utf-8 -*- +import json import re from elasticsearch_dsl import Q, Search, A from typing import List, Optional, Tuple, Dict, Union from dataclasses import dataclass + +from rag.settings import es_logger from rag.utils import rmSpace from rag.nlp import huqie, query import numpy as np @@ -34,30 +37,30 @@ class Dealer: group_docs: List[List] = None def _vector(self, txt, sim=0.8, topk=10): + qv = self.emb_mdl.encode_queries(txt) return { - "field": "q_vec", + "field": "q_%d_vec"%len(qv), "k": topk, "similarity": sim, "num_candidates": 1000, - "query_vector": self.emb_mdl.encode_queries(txt) + "query_vector": qv } 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")) + if req.get("doc_ids"): + bqry.filter.append(Q("terms", doc_id=req["doc_ids"])) 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"]) + src = req.get("fields", ["docnm_kwd", "content_ltks", "kb_id","img_id", + "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.highlight("content_ltks") @@ -66,22 +69,24 @@ class Dealer: 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=",./;:\\!()，。？：！……（）——、" - ) + 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"): s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps) s["knn"]["filter"] = bqry.to_dict() - del s["highlight"] + 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(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: bqry, _ = self.qryr.question(qst, min_match="10%") if req.get("kb_ids"): @@ -109,8 +114,7 @@ class Dealer: 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"]), + field=self.getFields(res, src), keywords=list(kwds) ) @@ -237,14 +241,4 @@ class Dealer: 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_*")) + diff --git a/rag/nlp/term_weight.py b/rag/nlp/term_weight.py index 14e8bfc3c..752bbf3ef 100644 --- a/rag/nlp/term_weight.py +++ b/rag/nlp/term_weight.py @@ -62,7 +62,7 @@ class Dealer: return set(res.keys()) 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 = {}, {} try: self.ne = json.load(open(os.path.join(fnm, "ner.json"), "r")) diff --git a/rag/settings.py b/rag/settings.py index 613abe75c..43deedf3d 100644 --- a/rag/settings.py +++ b/rag/settings.py @@ -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"); # you may not use this file except in compliance with the License. diff --git a/rag/svr/parse_user_docs.py b/rag/svr/parse_user_docs.py index 188662e6b..0de5d03dd 100644 --- a/rag/svr/parse_user_docs.py +++ b/rag/svr/parse_user_docs.py @@ -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"); # 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 # limitations under the License. # +import datetime import json import logging import os @@ -108,17 +109,17 @@ def build(row, cvmdl): (int(DOC_MAXIMUM_SIZE / 1024 / 1024))) return [] - res = ELASTICSEARCH.search(Q("term", doc_id=row["id"])) - if ELASTICSEARCH.getTotal(res) > 0: - ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]), - scripts=""" - if(!ctx._source.kb_id.contains('%s')) - ctx._source.kb_id.add('%s'); - """ % (str(row["kb_id"]), str(row["kb_id"])), - idxnm=search.index_name(row["tenant_id"]) - ) - set_progress(row["id"], 1, "Done") - return [] + # res = ELASTICSEARCH.search(Q("term", doc_id=row["id"])) + # if ELASTICSEARCH.getTotal(res) > 0: + # ELASTICSEARCH.updateScriptByQuery(Q("term", doc_id=row["id"]), + # scripts=""" + # if(!ctx._source.kb_id.contains('%s')) + # ctx._source.kb_id.add('%s'); + # """ % (str(row["kb_id"]), str(row["kb_id"])), + # idxnm=search.index_name(row["tenant_id"]) + # ) + # set_progress(row["id"], 1, "Done") + # return [] random.seed(time.time()) set_progress(row["id"], random.randint(0, 20) / @@ -155,8 +156,7 @@ def build(row, cvmdl): "doc_id": row["id"], "kb_id": [str(row["kb_id"])], "docnm_kwd": os.path.split(row["location"])[-1], - "title_tks": huqie.qie(row["name"]), - "updated_at": str(row["update_time"]).replace("T", " ")[:19] + "title_tks": huqie.qie(row["name"]) } doc["title_sm_tks"] = huqie.qieqie(doc["title_tks"]) output_buffer = BytesIO() @@ -179,6 +179,7 @@ def build(row, cvmdl): MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue()) d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"]) + d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19] docs.append(d) for arr, img in obj.table_chunks: @@ -193,6 +194,7 @@ def build(row, cvmdl): img.save(output_buffer, format='JPEG') MINIO.put(row["kb_id"], d["_id"], output_buffer.getvalue()) d["img_id"] = "{}-{}".format(row["kb_id"], d["_id"]) + d["create_time"] = str(datetime.datetime.now()).replace("T", " ")[:19] docs.append(d) set_progress(row["id"], random.randint(60, 70) / 100., "Continue embedding the content.") @@ -218,23 +220,11 @@ def embedding(docs, mdl): vects = 0.1 * tts + 0.9 * cnts assert len(vects) == len(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 -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): global model from rag.llm import HuEmbedding @@ -247,12 +237,12 @@ def main(comm, mod): tmf = open(tm_fnm, "a+") 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: set_progress(r["id"], -1, "Can't find embedding model!") cron_logger.error("Tenant({}) can't find embedding model!".format(r["tenant_id"])) continue - cv_mdl = model_instance(r["tenant_id"], LLMType.IMAGE2TEXT) + cv_mdl = TenantLLMService.model_instance(r["tenant_id"], LLMType.IMAGE2TEXT) st_tm = timer() cks = build(r, cv_mdl) if not cks: diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py index 632b01d6e..e036ca95e 100644 --- a/rag/utils/es_conn.py +++ b/rag/utils/es_conn.py @@ -241,6 +241,26 @@ class HuEs: es_logger.error("ES search timeout for 3 times!") 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): ubq = UpdateByQuery(index=self.idxnm).using(self.es).query(q) scripts = ""