71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
import os
|
|
import json
|
|
import requests
|
|
import numpy as np
|
|
from FlagEmbedding import FlagReranker
|
|
|
|
from src.config import RERANKER_LIST
|
|
from src.utils.logging_config import logger
|
|
|
|
|
|
class LocalReranker(FlagReranker):
|
|
def __init__(self, config, **kwargs):
|
|
model_info = RERANKER_LIST[config.reranker]
|
|
model_name_or_path = config.model_local_paths.get(model_info["name"], model_info.get("local_path"))
|
|
model_name_or_path = model_name_or_path or model_info["name"]
|
|
logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}")
|
|
|
|
super().__init__(model_name_or_path, use_fp16=True, device=config.device, **kwargs)
|
|
logger.info(f"Reranker model {config.reranker} loaded")
|
|
|
|
|
|
def sigmoid(x):
|
|
return 1 / (1 + np.exp(-x))
|
|
|
|
class SilconFlowReranker():
|
|
def __init__(self, config, **kwargs):
|
|
self.url = "https://api.siliconflow.cn/v1/rerank"
|
|
self.model = RERANKER_LIST[config.reranker]["name"]
|
|
|
|
api_key = os.getenv("SILICONFLOW_API_KEY")
|
|
assert api_key, "SILICONFLOW_API_KEY is required"
|
|
self.headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
def compute_score(self, sentence_pairs, batch_size = 256, max_length = 512, normalize = False):
|
|
# TODO 还没实现 batch_size
|
|
query, sentences = sentence_pairs[0], sentence_pairs[1]
|
|
payload = self.build_payload(query, sentences, max_length)
|
|
response = requests.request("POST", self.url, json=payload, headers=self.headers)
|
|
response = json.loads(response.text)
|
|
# logger.debug(f"SiliconFlow Reranker response: {response}")
|
|
|
|
results = sorted(response["results"], key=lambda x: x["index"])
|
|
all_scores = [result["relevance_score"] for result in results]
|
|
|
|
if normalize:
|
|
all_scores = [sigmoid(score) for score in all_scores]
|
|
|
|
return all_scores
|
|
|
|
def build_payload(self, query, sentences, max_length = 512):
|
|
return {
|
|
"model": self.model,
|
|
"query": query,
|
|
"documents": sentences,
|
|
"max_chunks_per_doc": max_length,
|
|
}
|
|
|
|
def get_reranker(config):
|
|
assert config.reranker in RERANKER_LIST.keys(), f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}"
|
|
provider, model_name = config.reranker.split('/', 1)
|
|
if provider == "local":
|
|
return LocalReranker(config)
|
|
elif provider == "siliconflow":
|
|
return SilconFlowReranker(config)
|
|
else:
|
|
raise ValueError(f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}")
|
|
|