feat: support xinference rerank model (#1466)

### What problem does this PR solve?

support xinference rerank model
#1455 

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
zhuhao
2024-07-11 18:37:41 +08:00
committed by GitHub
parent 9c023b6d8c
commit 009e18f094
4 changed files with 23 additions and 3 deletions

View File

@ -68,4 +68,5 @@ RerankModel = {
"BAAI": DefaultRerank,
"Jina": JinaRerank,
"Youdao": YoudaoRerank,
"Xinference": XInferenceRerank
}

View File

@ -136,4 +136,22 @@ class YoudaoRerank(DefaultRerank):
else: res.extend(scores)
return np.array(res), token_count
class XInferenceRerank(Base):
def __init__(self,model_name="",base_url=""):
self.model_name=model_name
self.base_url=base_url
self.headers = {
"Content-Type": "application/json",
"accept": "application/json"
}
def similarity(self, query: str, texts: list):
data = {
"model":self.model_name,
"query":query,
"return_documents": "true",
"return_len": "true",
"documents":texts
}
res = requests.post(self.base_url, headers=self.headers, json=data).json()
return np.array([d["relevance_score"] for d in res["results"]]),res["tokens"]["input_tokens"]+res["tokens"]["output_tokens"]