mirror of
https://github.com/infiniflow/ragflow.git
synced 2025-12-08 20:42:30 +08:00
Refa: make exception more clear. (#8224)
### What problem does this PR solve? #8156 ### Type of change - [x] Refactoring
This commit is contained in:
Kevin Hu
@ -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
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
Reference in New Issue
Block a user