pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 12:32:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

add dialog api (#33)

Browse Source
This commit is contained in:
KevinHuSh
2024-01-17 20:20:42 +08:00
committed by GitHub
parent 6be3dd56fa
commit 9bf75d4511
50 changed files with 511 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
rag/llm/__init__.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.

2
rag/llm/chat_model.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.

2
rag/llm/cv_model.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.

6
rag/llm/embedding_model.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.
@ -60,6 +60,10 @@ class HuEmbedding(Base):
res.extend(self.model.encode(texts[i:i + batch_size]).tolist())
return np.array(res), token_count
def encode_queries(self, text: str):
token_count = num_tokens_from_string(text)
return self.model.encode_queries([text]).tolist()[0], token_count
class OpenAIEmbed(Base):
def __init__(self, key, model_name="text-embedding-ada-002"):

2
rag/nlp/huqie.py
View File

@ -9,7 +9,7 @@ import string
import sys
from hanziconv import HanziConv
from web_server.utils.file_utils import get_project_base_directory
from api.utils.file_utils import get_project_base_directory
class Huqie:

2
rag/nlp/query.py
View File

@ -147,7 +147,7 @@ class EsQueryer:
atks = toDict(atks)
btkss = [toDict(tks) for tks in btkss]
tksim = [self.similarity(atks, btks) for btks in btkss]
return np.array(sims[0]) * vtweight + np.array(tksim) * tkweight
return np.array(sims[0]) * vtweight + np.array(tksim) * tkweight, sims[0], tksim
def similarity(self, qtwt, dtwt):
if isinstance(dtwt, type("")):

36
rag/nlp/search.py
View File

@ -15,7 +15,7 @@ def index_name(uid): return f"ragflow_{uid}"
class Dealer:
def __init__(self, es, emb_mdl):
def __init__(self, es):
self.qryr = query.EsQueryer(es)
self.qryr.flds = [
"title_tks^10",
@ -23,7 +23,6 @@ class Dealer:
"content_ltks^2",
"content_sm_ltks"]
self.es = es
self.emb_mdl = emb_mdl
@dataclass
class SearchResult:
@ -36,23 +35,26 @@ class Dealer:
keywords: Optional[List[str]] = None
group_docs: List[List] = None
def _vector(self, txt, sim=0.8, topk=10):
qv = self.emb_mdl.encode_queries(txt)
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": 1000,
"num_candidates": topk*2,
"query_vector": qv
}
def search(self, req, idxnm, tks_num=3):
def search(self, req, idxnm, emb_mdl=None):
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"]))
if req.get("doc_ids"):
bqry.filter.append(Q("terms", doc_id=req["doc_ids"]))
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})))
bqry.boost = 0.05
s = Search()
@ -60,7 +62,7 @@ class Dealer:
ps = int(req.get("size", 1000))
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"])
"q_1024_vec", "q_1536_vec", "available_int"])
s = s.query(bqry)[pg * ps:(pg + 1) * ps]
s = s.highlight("content_ltks")
@ -80,7 +82,8 @@ class Dealer:
s = s.to_dict()
q_vec = []
if req.get("vector"):
s["knn"] = self._vector(qst, req.get("similarity", 0.4), ps)
assert emb_mdl, "No embedding model selected"
s["knn"] = self._vector(qst, emb_mdl, req.get("similarity", 0.4), ps)
s["knn"]["filter"] = bqry.to_dict()
if "highlight" in s: del s["highlight"]
q_vec = s["knn"]["query_vector"]
@ -168,7 +171,7 @@ class Dealer:
def trans2floats(txt):
return [float(t) for t in txt.split("\t")]
def insert_citations(self, ans, top_idx, sres,
def insert_citations(self, ans, top_idx, sres, emb_mdl,
vfield="q_vec", cfield="content_ltks"):
ins_embd = [Dealer.trans2floats(
@ -179,15 +182,14 @@ class Dealer:
res = ""
def citeit():
nonlocal s, e, ans, res
nonlocal s, e, ans, res, emb_mdl
if not ins_embd:
return
embd = self.emb_mdl.encode(ans[s: e])
embd = 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
@ -225,20 +227,18 @@ class Dealer:
return res
def rerank(self, sres, query, tkweight=0.3, vtweight=0.7,
vfield="q_vec", cfield="content_ltks"):
def rerank(self, sres, query, tkweight=0.3, vtweight=0.7, cfield="content_ltks"):
ins_embd = [
Dealer.trans2floats(
sres.field[i]["q_vec"]) for i in sres.ids]
sres.field[i]["q_%d_vec"%len(sres.query_vector)]) 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,
sim, tksim, vtsim = self.qryr.hybrid_similarity(sres.query_vector,
ins_embd,
huqie.qie(query).split(" "),
ins_tw, tkweight, vtweight)
return sim
return sim, tksim, vtsim

2
rag/nlp/synonym.py
View File

@ -4,7 +4,7 @@ import time
import logging
import re
from web_server.utils.file_utils import get_project_base_directory
from api.utils.file_utils import get_project_base_directory
class Dealer:

2
rag/nlp/term_weight.py
View File

@ -5,7 +5,7 @@ import re
import os
import numpy as np
from rag.nlp import huqie
from web_server.utils.file_utils import get_project_base_directory
from api.utils.file_utils import get_project_base_directory
class Dealer:

8
rag/settings.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.
@ -14,9 +14,9 @@
# limitations under the License.
#
import os
from web_server.utils import get_base_config,decrypt_database_config
from web_server.utils.file_utils import get_project_base_directory
from web_server.utils.log_utils import LoggerFactory, getLogger
from api.utils import get_base_config,decrypt_database_config
from api.utils.file_utils import get_project_base_directory
from api.utils.log_utils import LoggerFactory, getLogger
# Server

15
rag/svr/parse_user_docs.py
View File

@ -1,5 +1,5 @@
#
# Copyright 2019 The RAG Flow Authors. All Rights Reserved.
# Copyright 2019 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.
@ -47,12 +47,12 @@ from rag.nlp.huchunk import (
PptChunker,
TextChunker
)
from web_server.db import LLMType
from web_server.db.services.document_service import DocumentService
from web_server.db.services.llm_service import TenantLLMService
from web_server.settings import database_logger
from web_server.utils import get_format_time
from web_server.utils.file_utils import get_project_base_directory
from api.db import LLMType
from api.db.services.document_service import DocumentService
from api.db.services.llm_service import TenantLLMService
from api.settings import database_logger
from api.utils import get_format_time
from api.utils.file_utils import get_project_base_directory
BATCH_SIZE = 64
@ -257,7 +257,6 @@ def main(comm, mod):
cron_logger.error(str(e))
continue
set_progress(r["id"], random.randint(70, 95) / 100.,
"Finished embedding! Start to build index!")
init_kb(r)

1
rag/utils/es_conn.py
View File

@ -66,7 +66,6 @@ class HuEs:
body=d,
id=id,
refresh=False,
doc_type="_doc",
retry_on_conflict=100)
es_logger.info("Successfully upsert: %s" % id)
T = True