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

add dockerfile for cuda envirement. Refine table search strategy, (#123)

Browse Source
This commit is contained in:
KevinHuSh
2024-03-14 19:45:29 +08:00
committed by GitHub
parent 937048e5fb
commit 675a9f8d9a
18 changed files with 259 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
rag/app/table.py
View File

@ -67,7 +67,7 @@ class Excel(ExcelParser):
def trans_datatime(s):
try:
return datetime_parse(s.strip()).strftime("%Y-%m-%dT%H:%M:%S")
return datetime_parse(s.strip()).strftime("%Y-%m-%d %H:%M:%S")
except Exception as e:
pass
@ -80,6 +80,7 @@ def trans_bool(s):
def column_data_type(arr):
arr = list(arr)
uni = len(set([a for a in arr if a is not None]))
counts = {"int": 0, "float": 0, "text": 0, "datetime": 0, "bool": 0}
trans = {t: f for f, t in
@ -130,7 +131,7 @@ def chunk(filename, binary=None, from_page=0, to_page=10000000000, lang="Chinese
if re.search(r"\.xlsx?$", filename, re.IGNORECASE):
callback(0.1, "Start to parse.")
excel_parser = Excel()
dfs = excel_parser(filename, binary, callback)
dfs = excel_parser(filename, binary, from_page=from_page, to_page=to_page, callback=callback)
elif re.search(r"\.(txt|csv)$", filename, re.IGNORECASE):
callback(0.1, "Start to parse.")
txt = ""
@ -188,7 +189,7 @@ def chunk(filename, binary=None, from_page=0, to_page=10000000000, lang="Chinese
df[clmns[j]] = cln
if ty == "text":
txts.extend([str(c) for c in cln if c])
clmns_map = [(py_clmns[i] + fieds_map[clmn_tys[i]], clmns[i])
clmns_map = [(py_clmns[i] + fieds_map[clmn_tys[i]], clmns[i].replace("_", " "))
for i in range(len(clmns))]
eng = lang.lower() == "english"#is_english(txts)
@ -201,6 +202,8 @@ def chunk(filename, binary=None, from_page=0, to_page=10000000000, lang="Chinese
for j in range(len(clmns)):
if row[clmns[j]] is None:
continue
if not str(row[clmns[j]]):
continue
fld = clmns_map[j][0]
d[fld] = row[clmns[j]] if clmn_tys[j] != "text" else huqie.qie(
row[clmns[j]])

13
rag/llm/__init__.py
View File

@ -19,18 +19,20 @@ from .cv_model import *
EmbeddingModel = {
"local": HuEmbedding,
"Local": HuEmbedding,
"OpenAI": OpenAIEmbed,
"通义千问": HuEmbedding, #QWenEmbed,
"智谱AI": ZhipuEmbed
"智谱AI": ZhipuEmbed,
"Moonshot": HuEmbedding
}
CvModel = {
"OpenAI": GptV4,
"local": LocalCV,
"Local": LocalCV,
"通义千问": QWenCV,
"智谱AI": Zhipu4V
"智谱AI": Zhipu4V,
"Moonshot": LocalCV
}
@ -38,6 +40,7 @@ ChatModel = {
"OpenAI": GptTurbo,
"智谱AI": ZhipuChat,
"通义千问": QWenChat,
"local": LocalLLM
"Local": LocalLLM,
"Moonshot": MoonshotChat
}

9
rag/llm/chat_model.py
View File

@ -14,11 +14,8 @@
# limitations under the License.
#
from abc import ABC
from copy import deepcopy
from openai import OpenAI
import openai
from rag.nlp import is_english
from rag.utils import num_tokens_from_string
@ -52,6 +49,12 @@ class GptTurbo(Base):
return "**ERROR**: "+str(e), 0
class MoonshotChat(GptTurbo):
def __init__(self, key, model_name="moonshot-v1-8k"):
self.client = OpenAI(api_key=key, base_url="https://api.moonshot.cn/v1",)
self.model_name = model_name
from dashscope import Generation
class QWenChat(Base):
def __init__(self, key, model_name=Generation.Models.qwen_turbo):

44
rag/llm/rpc_server.py
View File

@ -4,7 +4,7 @@ import random
import time
from multiprocessing.connection import Listener
from threading import Thread
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
class RPCHandler:
@ -47,14 +47,27 @@ tokenizer = None
def chat(messages, gen_conf):
global tokenizer
model = Model()
roles = {"system":"System", "user": "User", "assistant": "Assistant"}
line = ["{}: {}".format(roles[m["role"].lower()], m["content"]) for m in messages]
line = "\n".join(line) + "\nAssistant: "
tokens = tokenizer([line], return_tensors='pt')
tokens = {k: tokens[k].to(model.device) if isinstance(tokens[k], torch.Tensor) else tokens[k] for k in
tokens.keys()}
res = [tokenizer.decode(t) for t in model.generate(**tokens, **gen_conf)][0]
return res.split("Assistant: ")[-1]
try:
conf = {"max_new_tokens": int(gen_conf.get("max_tokens", 256)), "temperature": float(gen_conf.get("temperature", 0.1))}
print(messages, conf)
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
model_inputs.input_ids,
**conf
)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
return tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
except Exception as e:
return str(e)
def Model():
@ -71,20 +84,13 @@ if __name__ == "__main__":
handler = RPCHandler()
handler.register_function(chat)
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig
models = []
for _ in range(2):
for _ in range(1):
m = AutoModelForCausalLM.from_pretrained(args.model_name,
device_map="auto",
torch_dtype='auto',
trust_remote_code=True)
m.generation_config = GenerationConfig.from_pretrained(args.model_name)
m.generation_config.pad_token_id = m.generation_config.eos_token_id
torch_dtype='auto')
models.append(m)
tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=False,
trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
# Run the server
rpc_server(handler, ('0.0.0.0', args.port), authkey=b'infiniflow-token4kevinhu')

8
rag/nlp/search.py
View File

@ -7,6 +7,7 @@ from elasticsearch_dsl import Q, Search
from typing import List, Optional, Dict, Union
from dataclasses import dataclass
from api.settings import chat_logger
from rag.settings import es_logger
from rag.utils import rmSpace
from rag.nlp import huqie, query
@ -333,15 +334,16 @@ class Dealer:
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;fuzziness=AUTO:1,3;minimum_should_match=30%') ".format(fld, huqie.qieqie(huqie.qie(v)))
match = " MATCH({}, '{}', 'operator=OR;minimum_should_match=30%') ".format(fld, huqie.qieqie(huqie.qie(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)
es_logger.info(f"To es: {sql}")
chat_logger.info(f"To es: {sql}")
try:
tbl = self.es.sql(sql, fetch_size, format)
return tbl
except Exception as e:
es_logger.error(f"SQL failure: {sql} =>" + str(e))
chat_logger.error(f"SQL failure: {sql} =>" + str(e))
return {"error": str(e)}

17
rag/svr/task_executor.py
View File

@ -169,16 +169,25 @@ def init_kb(row):
def embedding(docs, mdl, parser_config={}, callback=None):
batch_size = 32
tts, cnts = [rmSpace(d["title_tks"]) for d in docs if d.get("title_tks")], [
d["content_with_weight"] for d in docs]
tk_count = 0
if len(tts) == len(cnts):
tts, c = mdl.encode(tts)
tk_count += c
tts_ = np.array([])
for i in range(0, len(tts), batch_size):
vts, c = mdl.encode(tts[i: i + batch_size])
if len(tts_) == 0:
tts_ = vts
else:
tts_ = np.concatenate((tts_, vts), axis=0)
tk_count += c
callback(prog=0.6 + 0.1 * (i + 1) / len(tts), msg="")
tts = tts_
cnts_ = np.array([])
for i in range(0, len(cnts), 8):
vts, c = mdl.encode(cnts[i: i+8])
for i in range(0, len(cnts), batch_size):
vts, c = mdl.encode(cnts[i: i+batch_size])
if len(cnts_) == 0: cnts_ = vts
else: cnts_ = np.concatenate((cnts_, vts), axis=0)
tk_count += c

2
rag/utils/es_conn.py
View File

@ -249,6 +249,8 @@ class HuEs:
except ConnectionTimeout as e:
es_logger.error("Timeout【Q】：" + sql)
continue
except Exception as e:
raise e
es_logger.error("ES search timeout for 3 times!")
raise ConnectionTimeout()