Skip to content

Commit

Permalink
feat: support siliconflow rerank (#7337)
Browse files Browse the repository at this point in the history
  • Loading branch information
Weaxs authored Aug 16, 2024
1 parent 7d4a0a4 commit 3a33062
Show file tree
Hide file tree
Showing 6 changed files with 146 additions and 1 deletion.
Empty file.
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
model: netease-youdao/bce-reranker-base_v1
model_type: rerank
model_properties:
context_size: 512
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
model: BAAI/bge-reranker-v2-m3
model_type: rerank
model_properties:
context_size: 8192
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
from typing import Optional

import httpx

from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
from core.model_runtime.errors.invoke import (
InvokeAuthorizationError,
InvokeBadRequestError,
InvokeConnectionError,
InvokeError,
InvokeRateLimitError,
InvokeServerUnavailableError,
)
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.__base.rerank_model import RerankModel


class SiliconflowRerankModel(RerankModel):

def _invoke(self, model: str, credentials: dict, query: str, docs: list[str],
score_threshold: Optional[float] = None, top_n: Optional[int] = None,
user: Optional[str] = None) -> RerankResult:
if len(docs) == 0:
return RerankResult(model=model, docs=[])

base_url = credentials.get('base_url', 'https://api.siliconflow.cn/v1')
if base_url.endswith('/'):
base_url = base_url[:-1]
try:
response = httpx.post(
base_url + '/rerank',
json={
"model": model,
"query": query,
"documents": docs,
"top_n": top_n,
"return_documents": True
},
headers={"Authorization": f"Bearer {credentials.get('api_key')}"}
)
response.raise_for_status()
results = response.json()

rerank_documents = []
for result in results['results']:
rerank_document = RerankDocument(
index=result['index'],
text=result['document']['text'],
score=result['relevance_score'],
)
if score_threshold is None or result['relevance_score'] >= score_threshold:
rerank_documents.append(rerank_document)

return RerankResult(model=model, docs=rerank_documents)
except httpx.HTTPStatusError as e:
raise InvokeServerUnavailableError(str(e))

def validate_credentials(self, model: str, credentials: dict) -> None:
try:

self._invoke(
model=model,
credentials=credentials,
query="What is the capital of the United States?",
docs=[
"Carson City is the capital city of the American state of Nevada. At the 2010 United States "
"Census, Carson City had a population of 55,274.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
"are a political division controlled by the United States. Its capital is Saipan.",
],
score_threshold=0.8
)
except Exception as ex:
raise CredentialsValidateFailedError(str(ex))

@property
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
"""
Map model invoke error to unified error
"""
return {
InvokeConnectionError: [httpx.ConnectError],
InvokeServerUnavailableError: [httpx.RemoteProtocolError],
InvokeRateLimitError: [],
InvokeAuthorizationError: [httpx.HTTPStatusError],
InvokeBadRequestError: [httpx.RequestError]
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,10 +12,11 @@ help:
en_US: Get your API Key from SiliconFlow
zh_Hans: 从 SiliconFlow 获取 API Key
url:
en_US: https://cloud.siliconflow.cn/keys
en_US: https://cloud.siliconflow.cn/account/ak
supported_model_types:
- llm
- text-embedding
- rerank
- speech2text
configurate_methods:
- predefined-model
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
import os

import pytest

from core.model_runtime.entities.rerank_entities import RerankResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.siliconflow.rerank.rerank import SiliconflowRerankModel


def test_validate_credentials():
model = SiliconflowRerankModel()

with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="BAAI/bge-reranker-v2-m3",
credentials={
"api_key": "invalid_key"
},
)

model.validate_credentials(
model="BAAI/bge-reranker-v2-m3",
credentials={
"api_key": os.environ.get("API_KEY"),
},
)


def test_invoke_model():
model = SiliconflowRerankModel()

result = model.invoke(
model='BAAI/bge-reranker-v2-m3',
credentials={
"api_key": os.environ.get("API_KEY"),
},
query="Who is Kasumi?",
docs=[
"Kasumi is a girl's name of Japanese origin meaning \"mist\".",
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
"and she leads a team named PopiParty."
],
score_threshold=0.8
)

assert isinstance(result, RerankResult)
assert len(result.docs) == 1
assert result.docs[0].index == 0
assert result.docs[0].score >= 0.8

0 comments on commit 3a33062

Please sign in to comment.