DifyでRAG(Text Embedding + Rerank)をセルフホスティングしたい! - Rerankモデルプロバイダーを用意する
February 05, 2025
Rerank用モデルプロバイダーの実装
Rerank用モデルは、https://huggingface.co/cl-nagoya/ruri-reranker-largeを用いる。
Rerank用モデルプロバイダーも、Text Embedding用モデルプロバイダーと同様の制約がある。 したがって、Text Embedding用モデルプロバイダーと同様に、Rerank用モデルプロバイダーを実装する。
実行環境と依存ライブラリの用意
依存ライブラリはText Embedding用モデルプロバイダーと同様である。 同じディレクトリ内で作業するとそのまま利用できる。
APIサーバーの実装
Rerank APIサーバーを、rerank-api-server.pyとして、次のように実装する。
from fastapi import FastAPI
from pydantic import BaseModel
from typing import List
from sentence_transformers import CrossEncoder
app = FastAPI()
reranker = CrossEncoder("cl-nagoya/ruri-reranker-large")
class ReRankRequest(BaseModel):
model: str
query: str
documents: List[str]
@app.post("/v1/rerank")
def rerank(request: ReRankRequest):
pairs = [(request.query, doc) for doc in request.documents]
scores = reranker.predict(pairs)
ranked = sorted(
zip(request.documents, scores, range(len(scores))),
key=lambda x: x[1],
reverse=True,
)
results = []
for doc, score, idx in ranked:
results.append(
{
"index": idx,
"document": {"text": doc},
"relevance_score": float(score),
}
)
res = {
"results": results,
"model": request.model,
"usage": {"total_tokens": 0},
}
return res
APIサーバーの起動
uvicornを通じて、APIサーバーを起動する。この際のポート番号はてきとうに空いているものを利用する。
当然ながら、Difyのインスタンスからネットワーク的に到達できるように、適宜変更が必要である。
$ uv run uvicorn rerank-api-server:app --host 0.0.0.0 --port 8082
useageの値は、トークンをいくら処理したかの値であり、主に課金額の計算などに用いられる。
オンプレで雑に運用するのであれば、0を返しても差し支えないと思われる。
動作検証
実際に機能するか検証する。 curlで試す際は、次のようにリクエストを投げる。
$ curl -v http://127.0.0.1:8082/v1/rerank -H 'Content-Type: application/json' --data-raw '
{
"model": "cl-nagoya/ruri-reranker-large",
"query": "山形県の蔵王温泉にある「泉質」はなに？",
"documents": [
"蔵王温泉はどのような特徴を 持つ温泉ですか？",
"山形市の蔵王温泉はどのような温泉ですか？",
"蔵王温泉の特徴は何ですか？"
]
}'
次のようなレスポンスが返ってくれば成功である。
$ curl -v http://127.0.0.1:8082/v1/rerank -H 'Content-Type: application/json' -d '
{
"model": "cl-nagoya/ruri-reranker-large",
"query": "山形県の蔵王温泉にある「泉質」はなに？",
"documents": [
"蔵王温泉はどのような特徴を 持つ温泉ですか？",
"山形市の蔵王温泉はどのような温泉ですか？",
"蔵王温泉の特徴は何ですか？"
]
}' | jq
% Total % Received % Xferd Average Speed Time Time Time Current
Dload Upload Total Spent Left Speed
0 0 0 0 0 0 0 0 --:--:-- --:--:-- --:--:-- 0* Trying 127.0.0.1:8082...
* Connected to 127.0.0.1 (127.0.0.1) port 8082
> POST /v1/rerank HTTP/1.1
> Host: 127.0.0.1:8082
> User-Agent: curl/8.5.0
> Accept: */*
> Content-Type: application/json
> Content-Length: 347
>
} [347 bytes data]
< HTTP/1.1 200 OK
< date: Thu, 06 Feb 2025 08:39:08 GMT
< server: uvicorn
< content-length: 465
< content-type: application/json
<
{ [465 bytes data]
100 812 100 465 100 347 3765 2810 --:--:-- --:--:-- --:--:-- 6601
* Connection #0 to host 127.0.0.1 left intact
{
"results": [
{
"document": {
"text": "蔵王温泉はどのような特徴を 持つ温泉ですか？"
},
"relevance_score": 0.029905224218964577,
"index": 0
},
{
"document": {
"text": "山形市の蔵王温泉はどのような温泉ですか？"
},
"relevance_score": 0.013406982645392418,
"index": 1
},
{
"document": {
"text": "蔵王温泉の特徴は何ですか？"
},
"relevance_score": 0.012443745508790016,
"index": 2
}
],
"model": "cl-nagoya/ruri-reranker-large",
"usage": {
"total_tokens": 0
}
}
indexの値が小さくrelevance_scoreの値が大きいものほど、より適切な内容となっているはずである。