OpenSource
/
ragflow
关注 4
点赞 0
派生 0
代码 工单 0 合并请求 0 版本发布 33 百科 动态
浏览代码

Refa: make exception more clear. (#8224) 

### What problem does this PR solve?

#8156

### Type of change
- [x] Refactoring
tags/v0.19.1
Kevin Hu 4 个月前
父节点
86a1411b07
当前提交
d36c8d18b1
没有帐户链接到提交者的电子邮件
共有 3 个文件被更改，包括 210 次插入 和 121 次删除
分列视图 显示文件统计
  1. 8
    1
      api/utils/log_utils.py
  2. 145
    88
      rag/llm/embedding_model.py
  3. 57
    32
      rag/llm/rerank_model.py

+ 8
- 1
api/utils/log_utils.py 查看文件

@@ -77,4 +77,11 @@ def initRootLogger(logfile_basename: str, log_format: str = "%(asctime)-15s %(le
pkg_logger.setLevel(pkg_level)

msg = f"{logfile_basename} log path: {log_path}, log levels: {pkg_levels}"
logger.info(msg)
logger.info(msg)


def log_exception(e, *args):
logging.exception(e)
for a in args:
logging.error(str(a))
raise e

+ 145
- 88
rag/llm/embedding_model.py 查看文件

@@ -19,7 +19,6 @@ import threading
from urllib.parse import urljoin

import requests
from requests.exceptions import JSONDecodeError
from huggingface_hub import snapshot_download
from zhipuai import ZhipuAI
import os
@@ -32,6 +31,7 @@ import asyncio

from api import settings
from api.utils.file_utils import get_home_cache_dir
from api.utils.log_utils import log_exception
from rag.utils import num_tokens_from_string, truncate
import google.generativeai as genai
import json
@@ -130,8 +130,11 @@ class OpenAIEmbed(Base):
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size],
model=self.model_name)
ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res)
try:
ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)
return np.array(ress), total_tokens

def encode_queries(self, text):
@@ -153,7 +156,10 @@ class LocalAIEmbed(Base):
ress = []
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data])
try:
ress.extend([d.embedding for d in res.data])
except Exception as _e:
log_exception(_e, res)
# local embedding for LmStudio donot count tokens
return np.array(ress), 1024

@@ -188,40 +194,39 @@ class QWenEmbed(Base):
def encode(self, texts: list):
import dashscope
batch_size = 4
try:
res = []
token_count = 0
texts = [truncate(t, 2048) for t in texts]
for i in range(0, len(texts), batch_size):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=texts[i:i + batch_size],
api_key=self.key,
text_type="document"
)
res = []
token_count = 0
texts = [truncate(t, 2048) for t in texts]
for i in range(0, len(texts), batch_size):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=texts[i:i + batch_size],
api_key=self.key,
text_type="document"
)
try:
embds = [[] for _ in range(len(resp["output"]["embeddings"]))]
for e in resp["output"]["embeddings"]:
embds[e["text_index"]] = e["embedding"]
res.extend(embds)
token_count += self.total_token_count(resp)
return np.array(res), token_count
except Exception as e:
raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name)
return np.array([]), 0
except Exception as _e:
log_exception(_e, resp)
raise
return np.array(res), token_count

def encode_queries(self, text):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=text[:2048],
api_key=self.key,
text_type="query"
)
try:
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=text[:2048],
api_key=self.key,
text_type="query"
)
return np.array(resp["output"]["embeddings"][0]
["embedding"]), self.total_token_count(resp)
except Exception:
raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name)
return np.array([]), 0
except Exception as _e:
log_exception(_e, resp)


class ZhipuEmbed(Base):
@@ -243,14 +248,20 @@ class ZhipuEmbed(Base):
for txt in texts:
res = self.client.embeddings.create(input=txt,
model=self.model_name)
arr.append(res.data[0].embedding)
tks_num += self.total_token_count(res)
try:
arr.append(res.data[0].embedding)
tks_num += self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)
return np.array(arr), tks_num

def encode_queries(self, text):
res = self.client.embeddings.create(input=text,
model=self.model_name)
return np.array(res.data[0].embedding), self.total_token_count(res)
try:
return np.array(res.data[0].embedding), self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)


class OllamaEmbed(Base):
@@ -266,7 +277,10 @@ class OllamaEmbed(Base):
res = self.client.embeddings(prompt=txt,
model=self.model_name,
options={"use_mmap": True})
arr.append(res["embedding"])
try:
arr.append(res["embedding"])
except Exception as _e:
log_exception(_e, res)
tks_num += 128
return np.array(arr), tks_num

@@ -274,7 +288,10 @@ class OllamaEmbed(Base):
res = self.client.embeddings(prompt=text,
model=self.model_name,
options={"use_mmap": True})
return np.array(res["embedding"]), 128
try:
return np.array(res["embedding"]), 128
except Exception as _e:
log_exception(_e, res)


class FastEmbed(DefaultEmbedding):
@@ -334,14 +351,20 @@ class XinferenceEmbed(Base):
total_tokens = 0
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res)
try:
ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)
return np.array(ress), total_tokens

def encode_queries(self, text):
res = self.client.embeddings.create(input=[text],
model=self.model_name)
return np.array(res.data[0].embedding), self.total_token_count(res)
try:
return np.array(res.data[0].embedding), self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)


class YoudaoEmbed(Base):
@@ -401,11 +424,10 @@ class JinaEmbed(Base):
response = requests.post(self.base_url, headers=self.headers, json=data)
try:
res = response.json()
except JSONDecodeError as e:
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}")
raise
ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res)
ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res)
except Exception as _e:
log_exception(_e, response)
return np.array(ress), token_count

def encode_queries(self, text):
@@ -468,14 +490,20 @@ class MistralEmbed(Base):
for i in range(0, len(texts), batch_size):
res = self.client.embeddings(input=texts[i:i + batch_size],
model=self.model_name)
ress.extend([d.embedding for d in res.data])
token_count += self.total_token_count(res)
try:
ress.extend([d.embedding for d in res.data])
token_count += self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)
return np.array(ress), token_count

def encode_queries(self, text):
res = self.client.embeddings(input=[truncate(text, 8196)],
model=self.model_name)
return np.array(res.data[0].embedding), self.total_token_count(res)
try:
return np.array(res.data[0].embedding), self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)


class BedrockEmbed(Base):
@@ -505,9 +533,12 @@ class BedrockEmbed(Base):
body = {"texts": [text], "input_type": 'search_document'}

response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read())
embeddings.extend([model_response["embedding"]])
token_count += num_tokens_from_string(text)
try:
model_response = json.loads(response["body"].read())
embeddings.extend([model_response["embedding"]])
token_count += num_tokens_from_string(text)
except Exception as _e:
log_exception(_e, response)

return np.array(embeddings), token_count

@@ -520,8 +551,11 @@ class BedrockEmbed(Base):
body = {"texts": [truncate(text, 8196)], "input_type": 'search_query'}

response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read())
embeddings.extend(model_response["embedding"])
try:
model_response = json.loads(response["body"].read())
embeddings.extend(model_response["embedding"])
except Exception as _e:
log_exception(_e, response)

return np.array(embeddings), token_count

@@ -544,7 +578,10 @@ class GeminiEmbed(Base):
content=texts[i: i + batch_size],
task_type="retrieval_document",
title="Embedding of single string")
ress.extend(result['embedding'])
try:
ress.extend(result['embedding'])
except Exception as _e:
log_exception(_e, result)
return np.array(ress),token_count
def encode_queries(self, text):
@@ -555,7 +592,10 @@ class GeminiEmbed(Base):
task_type="retrieval_document",
title="Embedding of single string")
token_count = num_tokens_from_string(text)
return np.array(result['embedding']), token_count
try:
return np.array(result['embedding']), token_count
except Exception as _e:
log_exception(_e, result)


class NvidiaEmbed(Base):
@@ -593,9 +633,8 @@ class NvidiaEmbed(Base):
response = requests.post(self.base_url, headers=self.headers, json=payload)
try:
res = response.json()
except JSONDecodeError as e:
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}")
raise
except Exception as _e:
log_exception(_e, response)
ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res)
return np.array(ress), token_count
@@ -641,8 +680,11 @@ class CoHereEmbed(Base):
input_type="search_document",
embedding_types=["float"],
)
ress.extend([d for d in res.embeddings.float])
token_count += res.meta.billed_units.input_tokens
try:
ress.extend([d for d in res.embeddings.float])
token_count += res.meta.billed_units.input_tokens
except Exception as _e:
log_exception(_e, res)
return np.array(ress), token_count

def encode_queries(self, text):
@@ -652,9 +694,10 @@ class CoHereEmbed(Base):
input_type="search_query",
embedding_types=["float"],
)
return np.array(res.embeddings.float[0]), int(
res.meta.billed_units.input_tokens
)
try:
return np.array(res.embeddings.float[0]), int(res.meta.billed_units.input_tokens)
except Exception as _e:
log_exception(_e, res)


class TogetherAIEmbed(OpenAIEmbed):
@@ -706,13 +749,11 @@ class SILICONFLOWEmbed(Base):
response = requests.post(self.base_url, json=payload, headers=self.headers)
try:
res = response.json()
except JSONDecodeError as e:
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}")
raise
if "data" not in res or not isinstance(res["data"], list) or len(res["data"]) != len(texts_batch):
raise ValueError(f"SILICONFLOWEmbed.encode got invalid response from {self.base_url}")
ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res)
ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res)
except Exception as _e:
log_exception(_e, response)

return np.array(ress), token_count

def encode_queries(self, text):
@@ -724,12 +765,9 @@ class SILICONFLOWEmbed(Base):
response = requests.post(self.base_url, json=payload, headers=self.headers).json()
try:
res = response.json()
except JSONDecodeError as e:
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}")
raise
if "data" not in res or not isinstance(res["data"], list) or len(res["data"])!= 1:
raise ValueError(f"SILICONFLOWEmbed.encode_queries got invalid response from {self.base_url}")
return np.array(res["data"][0]["embedding"]), self.total_token_count(res)
return np.array(res["data"][0]["embedding"]), self.total_token_count(res)
except Exception as _e:
log_exception(_e, response)


class ReplicateEmbed(Base):
@@ -765,17 +803,23 @@ class BaiduYiyanEmbed(Base):

def encode(self, texts: list, batch_size=16):
res = self.client.do(model=self.model_name, texts=texts).body
return (
np.array([r["embedding"] for r in res["data"]]),
self.total_token_count(res),
)
try:
return (
np.array([r["embedding"] for r in res["data"]]),
self.total_token_count(res),
)
except Exception as _e:
log_exception(_e, res)

def encode_queries(self, text):
res = self.client.do(model=self.model_name, texts=[text]).body
return (
np.array([r["embedding"] for r in res["data"]]),
self.total_token_count(res),
)
try:
return (
np.array([r["embedding"] for r in res["data"]]),
self.total_token_count(res),
)
except Exception as _e:
log_exception(_e, res)


class VoyageEmbed(Base):
@@ -793,15 +837,21 @@ class VoyageEmbed(Base):
res = self.client.embed(
texts=texts[i : i + batch_size], model=self.model_name, input_type="document"
)
ress.extend(res.embeddings)
token_count += res.total_tokens
try:
ress.extend(res.embeddings)
token_count += res.total_tokens
except Exception as _e:
log_exception(_e, res)
return np.array(ress), token_count

def encode_queries(self, text):
res = self.client.embed(
texts=text, model=self.model_name, input_type="query"
)
return np.array(res.embeddings)[0], res.total_tokens
try:
return np.array(res.embeddings)[0], res.total_tokens
except Exception as _e:
log_exception(_e, res)


class HuggingFaceEmbed(Base):
@@ -821,11 +871,14 @@ class HuggingFaceEmbed(Base):
headers={'Content-Type': 'application/json'}
)
if response.status_code == 200:
embedding = response.json()
embeddings.append(embedding[0])
try:
embedding = response.json()
embeddings.append(embedding[0])
return np.array(embeddings), sum([num_tokens_from_string(text) for text in texts])
except Exception as _e:
log_exception(_e, response)
else:
raise Exception(f"Error: {response.status_code} - {response.text}")
return np.array(embeddings), sum([num_tokens_from_string(text) for text in texts])

def encode_queries(self, text):
response = requests.post(
@@ -834,8 +887,11 @@ class HuggingFaceEmbed(Base):
headers={'Content-Type': 'application/json'}
)
if response.status_code == 200:
embedding = response.json()
return np.array(embedding[0]), num_tokens_from_string(text)
try:
embedding = response.json()
return np.array(embedding[0]), num_tokens_from_string(text)
except Exception as _e:
log_exception(_e, response)
else:
raise Exception(f"Error: {response.status_code} - {response.text}")

@@ -848,6 +904,7 @@ class VolcEngineEmbed(OpenAIEmbed):
model_name = json.loads(key).get('ep_id', '') + json.loads(key).get('endpoint_id', '')
super().__init__(ark_api_key,model_name,base_url)


class GPUStackEmbed(OpenAIEmbed):
def __init__(self, key, model_name, base_url):
if not base_url:

+ 57
- 32
rag/llm/rerank_model.py 查看文件

@@ -28,6 +28,7 @@ from yarl import URL

from api import settings
from api.utils.file_utils import get_home_cache_dir
from api.utils.log_utils import log_exception
from rag.utils import num_tokens_from_string, truncate
import json

@@ -170,8 +171,11 @@ class JinaRerank(Base):
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float)
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
return rank, self.total_token_count(res)


@@ -238,8 +242,11 @@ class XInferenceRerank(Base):
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float)
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
return rank, token_count


@@ -269,10 +276,11 @@ class LocalAIRerank(Base):
token_count += num_tokens_from_string(t)
res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float)
if 'results' not in res:
raise ValueError("response not contains results\n" + str(res))
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)

# Normalize the rank values to the range 0 to 1
min_rank = np.min(rank)
@@ -322,8 +330,11 @@ class NvidiaRerank(Base):
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float)
for d in res["rankings"]:
rank[d["index"]] = d["logit"]
try:
for d in res["rankings"]:
rank[d["index"]] = d["logit"]
except Exception as _e:
log_exception(_e, res)
return rank, token_count


@@ -361,10 +372,11 @@ class OpenAI_APIRerank(Base):
token_count += num_tokens_from_string(t)
res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float)
if 'results' not in res:
raise ValueError("response not contains results\n" + str(res))
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)

# Normalize the rank values to the range 0 to 1
min_rank = np.min(rank)
@@ -398,8 +410,11 @@ class CoHereRerank(Base):
return_documents=False,
)
rank = np.zeros(len(texts), dtype=float)
for d in res.results:
rank[d.index] = d.relevance_score
try:
for d in res.results:
rank[d.index] = d.relevance_score
except Exception as _e:
log_exception(_e, res)
return rank, token_count


@@ -439,11 +454,11 @@ class SILICONFLOWRerank(Base):
self.base_url, json=payload, headers=self.headers
).json()
rank = np.zeros(len(texts), dtype=float)
if "results" not in response:
return rank, 0
for d in response["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in response["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, response)
return (
rank,
response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"],
@@ -468,8 +483,11 @@ class BaiduYiyanRerank(Base):
top_n=len(texts),
).body
rank = np.zeros(len(texts), dtype=float)
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
try:
for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
return rank, self.total_token_count(res)


@@ -487,8 +505,11 @@ class VoyageRerank(Base):
res = self.client.rerank(
query=query, documents=texts, model=self.model_name, top_k=len(texts)
)
for r in res.results:
rank[r.index] = r.relevance_score
try:
for r in res.results:
rank[r.index] = r.relevance_score
except Exception as _e:
log_exception(_e, res)
return rank, res.total_tokens


@@ -511,8 +532,11 @@ class QWenRerank(Base):
)
rank = np.zeros(len(texts), dtype=float)
if resp.status_code == HTTPStatus.OK:
for r in resp.output.results:
rank[r.index] = r.relevance_score
try:
for r in resp.output.results:
rank[r.index] = r.relevance_score
except Exception as _e:
log_exception(_e, resp)
return rank, resp.usage.total_tokens
else:
raise ValueError(f"Error calling QWenRerank model {self.model_name}: {resp.status_code} - {resp.text}")
@@ -529,6 +553,7 @@ class HuggingfaceRerank(DefaultRerank):
res = requests.post(f"http://{url}/rerank", headers={"Content-Type": "application/json"},
json={"query": query, "texts": texts[i: i + batch_size],
"raw_scores": False, "truncate": True})

for o in res.json():
scores[o["index"] + i] = o["score"]
except Exception as e:
@@ -582,15 +607,15 @@ class GPUStackRerank(Base):
response_json = response.json()

rank = np.zeros(len(texts), dtype=float)
if "results" not in response_json:
return rank, 0

token_count = 0
for t in texts:
token_count += num_tokens_from_string(t)

for result in response_json["results"]:
rank[result["index"]] = result["relevance_score"]
try:
for result in response_json["results"]:
rank[result["index"]] = result["relevance_score"]
except Exception as _e:
log_exception(_e, response)

return (
rank,

撰写 预览
正在加载...
取消
保存