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

Refa: make exception more clear. (#8224)

Browse Source 
### What problem does this PR solve?

#8156

### Type of change
- [x] Refactoring
This commit is contained in:
Kevin Hu
2025-06-12 17:53:59 +08:00
committed by GitHub
parent 86a1411b07
commit d36c8d18b1
3 changed files with 210 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
api/utils/log_utils.py
View File

@ -77,4 +77,11 @@ def initRootLogger(logfile_basename: str, log_format: str = "%(asctime)-15s %(le
pkg_logger.setLevel(pkg_level) pkg_logger.setLevel(pkg_level)
msg = f"{logfile_basename} log path: {log_path}, log levels: {pkg_levels}" 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

233
rag/llm/embedding_model.py
View File

@ -19,7 +19,6 @@ import threading
from urllib.parse import urljoin from urllib.parse import urljoin
import requests import requests
from requests.exceptions import JSONDecodeError
from huggingface_hub import snapshot_download from huggingface_hub import snapshot_download
from zhipuai import ZhipuAI from zhipuai import ZhipuAI
import os import os
@ -32,6 +31,7 @@ import asyncio
from api import settings from api import settings
from api.utils.file_utils import get_home_cache_dir 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 from rag.utils import num_tokens_from_string, truncate
import google.generativeai as genai import google.generativeai as genai
import json import json
@ -130,8 +130,11 @@ class OpenAIEmbed(Base):
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], res = self.client.embeddings.create(input=texts[i:i + batch_size],
model=self.model_name) model=self.model_name)
ress.extend([d.embedding for d in res.data]) try:
total_tokens += self.total_token_count(res) 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 return np.array(ress), total_tokens
def encode_queries(self, text): def encode_queries(self, text):
@ -153,7 +156,10 @@ class LocalAIEmbed(Base):
ress = [] ress = []
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name) 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 # local embedding for LmStudio donot count tokens
return np.array(ress), 1024 return np.array(ress), 1024
@ -188,40 +194,39 @@ class QWenEmbed(Base):
def encode(self, texts: list): def encode(self, texts: list):
import dashscope import dashscope
batch_size = 4 batch_size = 4
try: res = []
res = [] token_count = 0
token_count = 0 texts = [truncate(t, 2048) for t in texts]
texts = [truncate(t, 2048) for t in texts] for i in range(0, len(texts), batch_size):
for i in range(0, len(texts), batch_size): resp = dashscope.TextEmbedding.call(
resp = dashscope.TextEmbedding.call( model=self.model_name,
model=self.model_name, input=texts[i:i + batch_size],
input=texts[i:i + batch_size], api_key=self.key,
api_key=self.key, text_type="document"
text_type="document" )
) try:
embds = [[] for _ in range(len(resp["output"]["embeddings"]))] embds = [[] for _ in range(len(resp["output"]["embeddings"]))]
for e in resp["output"]["embeddings"]: for e in resp["output"]["embeddings"]:
embds[e["text_index"]] = e["embedding"] embds[e["text_index"]] = e["embedding"]
res.extend(embds) res.extend(embds)
token_count += self.total_token_count(resp) token_count += self.total_token_count(resp)
return np.array(res), token_count except Exception as _e:
except Exception as e: log_exception(_e, resp)
raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name) raise
return np.array([]), 0 return np.array(res), token_count
def encode_queries(self, text): def encode_queries(self, text):
resp = dashscope.TextEmbedding.call(
model=self.model_name,
input=text[:2048],
api_key=self.key,
text_type="query"
)
try: 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] return np.array(resp["output"]["embeddings"][0]
["embedding"]), self.total_token_count(resp) ["embedding"]), self.total_token_count(resp)
except Exception: except Exception as _e:
raise Exception("Account abnormal. Please ensure it's on good standing to use QWen's "+self.model_name) log_exception(_e, resp)
return np.array([]), 0
class ZhipuEmbed(Base): class ZhipuEmbed(Base):
@ -243,14 +248,20 @@ class ZhipuEmbed(Base):
for txt in texts: for txt in texts:
res = self.client.embeddings.create(input=txt, res = self.client.embeddings.create(input=txt,
model=self.model_name) model=self.model_name)
arr.append(res.data[0].embedding) try:
tks_num += self.total_token_count(res) 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 return np.array(arr), tks_num
def encode_queries(self, text): def encode_queries(self, text):
res = self.client.embeddings.create(input=text, res = self.client.embeddings.create(input=text,
model=self.model_name) 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): class OllamaEmbed(Base):
@ -266,7 +277,10 @@ class OllamaEmbed(Base):
res = self.client.embeddings(prompt=txt, res = self.client.embeddings(prompt=txt,
model=self.model_name, model=self.model_name,
options={"use_mmap": True}) options={"use_mmap": True})
arr.append(res["embedding"]) try:
arr.append(res["embedding"])
except Exception as _e:
log_exception(_e, res)
tks_num += 128 tks_num += 128
return np.array(arr), tks_num return np.array(arr), tks_num
@ -274,7 +288,10 @@ class OllamaEmbed(Base):
res = self.client.embeddings(prompt=text, res = self.client.embeddings(prompt=text,
model=self.model_name, model=self.model_name,
options={"use_mmap": True}) 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): class FastEmbed(DefaultEmbedding):
@ -334,14 +351,20 @@ class XinferenceEmbed(Base):
total_tokens = 0 total_tokens = 0
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i:i + batch_size], model=self.model_name) 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:
total_tokens += self.total_token_count(res) 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 return np.array(ress), total_tokens
def encode_queries(self, text): def encode_queries(self, text):
res = self.client.embeddings.create(input=[text], res = self.client.embeddings.create(input=[text],
model=self.model_name) 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): class YoudaoEmbed(Base):
@ -401,11 +424,10 @@ class JinaEmbed(Base):
response = requests.post(self.base_url, headers=self.headers, json=data) response = requests.post(self.base_url, headers=self.headers, json=data)
try: try:
res = response.json() res = response.json()
except JSONDecodeError as e: ress.extend([d["embedding"] for d in res["data"]])
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") token_count += self.total_token_count(res)
raise except Exception as _e:
ress.extend([d["embedding"] for d in res["data"]]) log_exception(_e, response)
token_count += self.total_token_count(res)
return np.array(ress), token_count return np.array(ress), token_count
def encode_queries(self, text): def encode_queries(self, text):
@ -468,14 +490,20 @@ class MistralEmbed(Base):
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
res = self.client.embeddings(input=texts[i:i + batch_size], res = self.client.embeddings(input=texts[i:i + batch_size],
model=self.model_name) model=self.model_name)
ress.extend([d.embedding for d in res.data]) try:
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, res)
return np.array(ress), token_count return np.array(ress), token_count
def encode_queries(self, text): def encode_queries(self, text):
res = self.client.embeddings(input=[truncate(text, 8196)], res = self.client.embeddings(input=[truncate(text, 8196)],
model=self.model_name) 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): class BedrockEmbed(Base):
@ -505,9 +533,12 @@ class BedrockEmbed(Base):
body = {"texts": [text], "input_type": 'search_document'} body = {"texts": [text], "input_type": 'search_document'}
response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body)) response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read()) try:
embeddings.extend([model_response["embedding"]]) model_response = json.loads(response["body"].read())
token_count += num_tokens_from_string(text) 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 return np.array(embeddings), token_count
@ -520,8 +551,11 @@ class BedrockEmbed(Base):
body = {"texts": [truncate(text, 8196)], "input_type": 'search_query'} body = {"texts": [truncate(text, 8196)], "input_type": 'search_query'}
response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body)) response = self.client.invoke_model(modelId=self.model_name, body=json.dumps(body))
model_response = json.loads(response["body"].read()) try:
embeddings.extend(model_response["embedding"]) 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 return np.array(embeddings), token_count
@ -544,7 +578,10 @@ class GeminiEmbed(Base):
content=texts[i: i + batch_size], content=texts[i: i + batch_size],
task_type="retrieval_document", task_type="retrieval_document",
title="Embedding of single string") 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 return np.array(ress),token_count
def encode_queries(self, text): def encode_queries(self, text):
@ -555,7 +592,10 @@ class GeminiEmbed(Base):
task_type="retrieval_document", task_type="retrieval_document",
title="Embedding of single string") title="Embedding of single string")
token_count = num_tokens_from_string(text) 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): class NvidiaEmbed(Base):
@ -593,9 +633,8 @@ class NvidiaEmbed(Base):
response = requests.post(self.base_url, headers=self.headers, json=payload) response = requests.post(self.base_url, headers=self.headers, json=payload)
try: try:
res = response.json() res = response.json()
except JSONDecodeError as e: except Exception as _e:
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") log_exception(_e, response)
raise
ress.extend([d["embedding"] for d in res["data"]]) ress.extend([d["embedding"] for d in res["data"]])
token_count += self.total_token_count(res) token_count += self.total_token_count(res)
return np.array(ress), token_count return np.array(ress), token_count
@ -641,8 +680,11 @@ class CoHereEmbed(Base):
input_type="search_document", input_type="search_document",
embedding_types=["float"], embedding_types=["float"],
) )
ress.extend([d for d in res.embeddings.float]) try:
token_count += res.meta.billed_units.input_tokens 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 return np.array(ress), token_count
def encode_queries(self, text): def encode_queries(self, text):
@ -652,9 +694,10 @@ class CoHereEmbed(Base):
input_type="search_query", input_type="search_query",
embedding_types=["float"], embedding_types=["float"],
) )
return np.array(res.embeddings.float[0]), int( try:
res.meta.billed_units.input_tokens 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): class TogetherAIEmbed(OpenAIEmbed):
@ -706,13 +749,11 @@ class SILICONFLOWEmbed(Base):
response = requests.post(self.base_url, json=payload, headers=self.headers) response = requests.post(self.base_url, json=payload, headers=self.headers)
try: try:
res = response.json() res = response.json()
except JSONDecodeError as e: ress.extend([d["embedding"] for d in res["data"]])
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") token_count += self.total_token_count(res)
raise except Exception as _e:
if "data" not in res or not isinstance(res["data"], list) or len(res["data"]) != len(texts_batch): log_exception(_e, response)
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)
return np.array(ress), token_count return np.array(ress), token_count
def encode_queries(self, text): def encode_queries(self, text):
@ -724,12 +765,9 @@ class SILICONFLOWEmbed(Base):
response = requests.post(self.base_url, json=payload, headers=self.headers).json() response = requests.post(self.base_url, json=payload, headers=self.headers).json()
try: try:
res = response.json() res = response.json()
except JSONDecodeError as e: return np.array(res["data"][0]["embedding"]), self.total_token_count(res)
logging.error(f"JSON decode error: {e}\nResponse content: {response.text[:2000]}") except Exception as _e:
raise log_exception(_e, response)
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)
class ReplicateEmbed(Base): class ReplicateEmbed(Base):
@ -765,17 +803,23 @@ class BaiduYiyanEmbed(Base):
def encode(self, texts: list, batch_size=16): def encode(self, texts: list, batch_size=16):
res = self.client.do(model=self.model_name, texts=texts).body res = self.client.do(model=self.model_name, texts=texts).body
return ( try:
np.array([r["embedding"] for r in res["data"]]), return (
self.total_token_count(res), 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): def encode_queries(self, text):
res = self.client.do(model=self.model_name, texts=[text]).body res = self.client.do(model=self.model_name, texts=[text]).body
return ( try:
np.array([r["embedding"] for r in res["data"]]), return (
self.total_token_count(res), 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): class VoyageEmbed(Base):
@ -793,15 +837,21 @@ class VoyageEmbed(Base):
res = self.client.embed( res = self.client.embed(
texts=texts[i : i + batch_size], model=self.model_name, input_type="document" texts=texts[i : i + batch_size], model=self.model_name, input_type="document"
) )
ress.extend(res.embeddings) try:
token_count += res.total_tokens ress.extend(res.embeddings)
token_count += res.total_tokens
except Exception as _e:
log_exception(_e, res)
return np.array(ress), token_count return np.array(ress), token_count
def encode_queries(self, text): def encode_queries(self, text):
res = self.client.embed( res = self.client.embed(
texts=text, model=self.model_name, input_type="query" 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): class HuggingFaceEmbed(Base):
@ -821,11 +871,14 @@ class HuggingFaceEmbed(Base):
headers={'Content-Type': 'application/json'} headers={'Content-Type': 'application/json'}
) )
if response.status_code == 200: if response.status_code == 200:
embedding = response.json() try:
embeddings.append(embedding[0]) 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: else:
raise Exception(f"Error: {response.status_code} - {response.text}") 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): def encode_queries(self, text):
response = requests.post( response = requests.post(
@ -834,8 +887,11 @@ class HuggingFaceEmbed(Base):
headers={'Content-Type': 'application/json'} headers={'Content-Type': 'application/json'}
) )
if response.status_code == 200: if response.status_code == 200:
embedding = response.json() try:
return np.array(embedding[0]), num_tokens_from_string(text) embedding = response.json()
return np.array(embedding[0]), num_tokens_from_string(text)
except Exception as _e:
log_exception(_e, response)
else: else:
raise Exception(f"Error: {response.status_code} - {response.text}") 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', '') model_name = json.loads(key).get('ep_id', '') + json.loads(key).get('endpoint_id', '')
super().__init__(ark_api_key,model_name,base_url) super().__init__(ark_api_key,model_name,base_url)
class GPUStackEmbed(OpenAIEmbed): class GPUStackEmbed(OpenAIEmbed):
def __init__(self, key, model_name, base_url): def __init__(self, key, model_name, base_url):
if not base_url: if not base_url:

89
rag/llm/rerank_model.py
View File

@ -28,6 +28,7 @@ from yarl import URL
from api import settings from api import settings
from api.utils.file_utils import get_home_cache_dir 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 from rag.utils import num_tokens_from_string, truncate
import json import json
@ -170,8 +171,11 @@ class JinaRerank(Base):
} }
res = requests.post(self.base_url, headers=self.headers, json=data).json() res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
for d in res["results"]: try:
rank[d["index"]] = d["relevance_score"] 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) 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() res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
for d in res["results"]: try:
rank[d["index"]] = d["relevance_score"] for d in res["results"]:
rank[d["index"]] = d["relevance_score"]
except Exception as _e:
log_exception(_e, res)
return rank, token_count return rank, token_count
@ -269,10 +276,11 @@ class LocalAIRerank(Base):
token_count += num_tokens_from_string(t) token_count += num_tokens_from_string(t)
res = requests.post(self.base_url, headers=self.headers, json=data).json() res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
if 'results' not in res: try:
raise ValueError("response not contains results\n" + str(res)) for d in res["results"]:
for d in res["results"]: rank[d["index"]] = d["relevance_score"]
rank[d["index"]] = d["relevance_score"] except Exception as _e:
log_exception(_e, res)
# Normalize the rank values to the range 0 to 1 # Normalize the rank values to the range 0 to 1
min_rank = np.min(rank) min_rank = np.min(rank)
@ -322,8 +330,11 @@ class NvidiaRerank(Base):
} }
res = requests.post(self.base_url, headers=self.headers, json=data).json() res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
for d in res["rankings"]: try:
rank[d["index"]] = d["logit"] for d in res["rankings"]:
rank[d["index"]] = d["logit"]
except Exception as _e:
log_exception(_e, res)
return rank, token_count return rank, token_count
@ -361,10 +372,11 @@ class OpenAI_APIRerank(Base):
token_count += num_tokens_from_string(t) token_count += num_tokens_from_string(t)
res = requests.post(self.base_url, headers=self.headers, json=data).json() res = requests.post(self.base_url, headers=self.headers, json=data).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
if 'results' not in res: try:
raise ValueError("response not contains results\n" + str(res)) for d in res["results"]:
for d in res["results"]: rank[d["index"]] = d["relevance_score"]
rank[d["index"]] = d["relevance_score"] except Exception as _e:
log_exception(_e, res)
# Normalize the rank values to the range 0 to 1 # Normalize the rank values to the range 0 to 1
min_rank = np.min(rank) min_rank = np.min(rank)
@ -398,8 +410,11 @@ class CoHereRerank(Base):
return_documents=False, return_documents=False,
) )
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
for d in res.results: try:
rank[d.index] = d.relevance_score for d in res.results:
rank[d.index] = d.relevance_score
except Exception as _e:
log_exception(_e, res)
return rank, token_count return rank, token_count
@ -439,11 +454,11 @@ class SILICONFLOWRerank(Base):
self.base_url, json=payload, headers=self.headers self.base_url, json=payload, headers=self.headers
).json() ).json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
if "results" not in response: try:
return rank, 0 for d in response["results"]:
rank[d["index"]] = d["relevance_score"]
for d in response["results"]: except Exception as _e:
rank[d["index"]] = d["relevance_score"] log_exception(_e, response)
return ( return (
rank, rank,
response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"], response["meta"]["tokens"]["input_tokens"] + response["meta"]["tokens"]["output_tokens"],
@ -468,8 +483,11 @@ class BaiduYiyanRerank(Base):
top_n=len(texts), top_n=len(texts),
).body ).body
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
for d in res["results"]: try:
rank[d["index"]] = d["relevance_score"] 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) return rank, self.total_token_count(res)
@ -487,8 +505,11 @@ class VoyageRerank(Base):
res = self.client.rerank( res = self.client.rerank(
query=query, documents=texts, model=self.model_name, top_k=len(texts) query=query, documents=texts, model=self.model_name, top_k=len(texts)
) )
for r in res.results: try:
rank[r.index] = r.relevance_score for r in res.results:
rank[r.index] = r.relevance_score
except Exception as _e:
log_exception(_e, res)
return rank, res.total_tokens return rank, res.total_tokens
@ -511,8 +532,11 @@ class QWenRerank(Base):
) )
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
if resp.status_code == HTTPStatus.OK: if resp.status_code == HTTPStatus.OK:
for r in resp.output.results: try:
rank[r.index] = r.relevance_score 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 return rank, resp.usage.total_tokens
else: else:
raise ValueError(f"Error calling QWenRerank model {self.model_name}: {resp.status_code} - {resp.text}") 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"}, res = requests.post(f"http://{url}/rerank", headers={"Content-Type": "application/json"},
json={"query": query, "texts": texts[i: i + batch_size], json={"query": query, "texts": texts[i: i + batch_size],
"raw_scores": False, "truncate": True}) "raw_scores": False, "truncate": True})
for o in res.json(): for o in res.json():
scores[o["index"] + i] = o["score"] scores[o["index"] + i] = o["score"]
except Exception as e: except Exception as e:
@ -582,15 +607,15 @@ class GPUStackRerank(Base):
response_json = response.json() response_json = response.json()
rank = np.zeros(len(texts), dtype=float) rank = np.zeros(len(texts), dtype=float)
if "results" not in response_json:
return rank, 0
token_count = 0 token_count = 0
for t in texts: for t in texts:
token_count += num_tokens_from_string(t) token_count += num_tokens_from_string(t)
try:
for result in response_json["results"]: for result in response_json["results"]:
rank[result["index"]] = result["relevance_score"] rank[result["index"]] = result["relevance_score"]
except Exception as _e:
log_exception(_e, response)
return ( return (
rank, rank,