|
|
|
@@ -136,10 +136,11 @@ 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 |
|
|
|
def __init__(self, key="xxxxxxx", model_name="", base_url=""): |
|
|
|
self.model_name = model_name |
|
|
|
self.base_url = base_url |
|
|
|
self.headers = { |
|
|
|
"Content-Type": "application/json", |
|
|
|
"accept": "application/json" |
|
|
|
@@ -147,11 +148,12 @@ class XInferenceRerank(Base): |
|
|
|
|
|
|
|
def similarity(self, query: str, texts: list): |
|
|
|
data = { |
|
|
|
"model":self.model_name, |
|
|
|
"query":query, |
|
|
|
"model": self.model_name, |
|
|
|
"query": query, |
|
|
|
"return_documents": "true", |
|
|
|
"return_len": "true", |
|
|
|
"documents":texts |
|
|
|
"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"] |
|
|
|
return np.array([d["relevance_score"] for d in res["results"]]), res["tokens"]["input_tokens"] + res["tokens"][ |
|
|
|
"output_tokens"] |