merge: bge-reranker service from GPU server

This commit is contained in:
yumoqing 2026-07-21 18:20:50 +08:00
commit 75e5bc85e6
8 changed files with 288 additions and 1 deletions

7
.gitignore vendored Normal file
View File

@ -0,0 +1,7 @@
ah.pid
nohup.out
__pycache__/
*.pyc
workers/__pycache__/
logs/
files/

View File

@ -1,2 +1,79 @@
# bge-reranker
# BGE-Reranker Service
BGE-Reranker-V2-M3 文本重排序服务,用于 RAG 检索后的精排阶段。
## 特性
- 模型: BAAI/bge-reranker-v2-m3 (2.2GB)
- 推理: FP16, GPU 加速
- 延迟: ~45ms (模型加载后)
- 最大文档数: 100条/请求
## 部署
```bash
cd /data/ymq/bge-reranker
bash build.sh deploy # 启动
bash build.sh stop # 停止
bash build.sh status # 状态
```
## API
### GET /api/status
服务状态。
### POST /api/rerank
```json
{
"query": "安装步骤",
"documents": ["打开包装箱", "拧下螺丝", "注意事项"],
"top_k": 3
}
```
返回:
```json
{
"status": "SUCCEEDED",
"query": "安装步骤",
"ranked_docs": [
{"doc": "拧下螺丝", "score": 0.164},
{"doc": "打开包装箱", "score": 0.065},
{"doc": "注意事项", "score": 0.017}
],
"total": 3,
"returned": 3,
"elapsed": 0.0456
}
```
## 在 RAG 中的角色
```
查询 → CLIP embedding → VDB 向量召回 top-20
→ FalkorDB 图扩展
→ 合并
→ BGE-Rerank (本服务) → top-5
→ LLM 生成答案
```
## 离线模型下载
```bash
# 使用 huggingface-cli
huggingface-cli download BAAI/bge-reranker-v2-m3 --local-dir /data/ymq/models/BAAI/bge-reranker-v2-m3
# 或使用 hf
hf download BAAI/bge-reranker-v2-m3 --local-dir /data/ymq/models/BAAI/bge-reranker-v2-m3
```
## 配置
- 端口: 9090
- GPU: 2
- 虚拟环境: /data/ymq/wan22-service/py3
## Git
git@git.opencomputing.cn:yumoqing/bge-reranker.git

6
ah.py Normal file
View File

@ -0,0 +1,6 @@
# -*- coding:utf-8 -*-
from ahserver.webapp import webapp
from init import load_bge_reranker
if __name__ == '__main__':
webapp(load_bge_reranker)

57
build.sh Executable file
View File

@ -0,0 +1,57 @@
#!/usr/bin/env bash
# BGE-Reranker Service
set -e
cd "$(dirname "$0")"
SERVICE_NAME="bge-reranker"
PORT=9090
GPU=2
PY=/data/ymq/wan22-service/py3/bin/python
action="${1:-status}"
case "$action" in
deploy|update)
echo "=== $SERVICE_NAME Deploy (GPU $GPU, port $PORT) ==="
if [ -f ah.pid ] && kill -0 $(cat ah.pid) 2>/dev/null; then
kill $(cat ah.pid) 2>/dev/null || true; sleep 2
fi
if [ -d .git ] && [ -f .git/HEAD ]; then
git pull origin master 2>/dev/null || true
fi
mkdir -p logs files wwwroot
export PYTHONPATH="$(pwd)"
export CUDA_VISIBLE_DEVICES=$GPU
nohup $PY ah.py > nohup.out 2>&1 &
echo $! > ah.pid
echo "Started PID $(cat ah.pid) on port $PORT (GPU $GPU)"
sleep 5
if curl -s http://localhost:$PORT/api/status > /dev/null 2>&1; then
echo "Service healthy"
else
echo "WARNING: not responding, check nohup.out"
tail -20 nohup.out
fi
;;
stop)
if [ -f ah.pid ]; then
kill $(cat ah.pid) 2>/dev/null || true; rm -f ah.pid; echo "Stopped"
else echo "Not running"; fi
;;
start)
mkdir -p logs files wwwroot
export PYTHONPATH="$(pwd)"; export CUDA_VISIBLE_DEVICES=$GPU
nohup $PY ah.py > nohup.out 2>&1 &
echo $! > ah.pid; echo "Started PID $(cat ah.pid)"
;;
status)
echo "=== $SERVICE_NAME Status ==="
if [ -f ah.pid ] && kill -0 $(cat ah.pid) 2>/dev/null; then
echo "Process: running (PID $(cat ah.pid))"
else echo "Process: not running"; fi
echo "Port: $PORT, GPU: $GPU"
if curl -s --max-time 3 http://localhost:$PORT/api/status > /dev/null 2>&1; then
echo "HTTP: OK"
else echo "HTTP: not responding"; fi
;;
*) echo "Usage: $0 {deploy|update|stop|start|status}"; exit 1 ;;
esac

25
conf/config.json Normal file
View File

@ -0,0 +1,25 @@
{
"password_key": "BgeReranker2026Key",
"filesroot": "$[workdir]$/files",
"logger": {
"name": "bge-reranker",
"levelname": "info",
"logfile": "$[workdir]$/logs/bge-reranker.log"
},
"website": {
"paths": [["$[workdir]$/wwwroot", ""]],
"client_max_size": 10485760,
"host": "0.0.0.0",
"port": 9090,
"coding": "utf-8",
"indexes": ["index.html"],
"startswiths": [
{"leading": "/api/status", "registerfunction": "status"},
{"leading": "/api/rerank", "registerfunction": "rerank"}
],
"processors": [
[".tmpl", "tmpl"], [".app", "app"], [".ui", "bui"],
[".dspy", "dspy"], [".md", "md"]
]
}
}

55
init.py Normal file
View File

@ -0,0 +1,55 @@
# -*- coding:utf-8 -*-
from traceback import format_exc
from ahserver.serverenv import ServerEnv
from appPublic.registerfunction import RegisterFunction
from appPublic.log import debug, exception
import json
async def status_handler(request, params_kw, *args, **kwargs):
"""Status endpoint"""
return json.dumps({
"service": "bge-reranker",
"model": "BAAI/bge-reranker-v2-m3",
"model_path": "/data/ymq/models/BAAI/bge-reranker-v2-m3",
"endpoints": ["/api/status", "/api/rerank"]
}, indent=2, ensure_ascii=False)
async def rerank_handler(request, params_kw, *args, **kwargs):
"""Rerank endpoint"""
try:
query = params_kw.get("query")
documents = params_kw.get("documents", [])
top_k = params_kw.get("top_k", 5)
if not query or not documents:
return json.dumps({"error": "query and documents required"})
if not isinstance(documents, list):
return json.dumps({"error": "documents must be a list"})
from workers.bge_reranker import rerank
import time
start = time.time()
result = rerank(query, documents, top_k)
elapsed = round(time.time() - start, 4)
return json.dumps({
"status": "SUCCEEDED",
"query": query,
"ranked_docs": result["ranked_docs"],
"total": result["total"],
"returned": result["returned"],
"elapsed": elapsed
}, ensure_ascii=False)
except Exception as e:
exception(f"{e}, {format_exc()}")
return json.dumps({"error": str(e)})
def load_bge_reranker():
"""Register API handlers"""
env = ServerEnv()
rf = RegisterFunction()
rf.register("status", status_handler)
rf.register("rerank", rerank_handler)

0
workers/__init__.py Normal file
View File

60
workers/bge_reranker.py Normal file
View File

@ -0,0 +1,60 @@
# -*- coding:utf-8 -*-
"""BGE-Reranker-M3 lazy-loading wrapper."""
import os
import json
import torch
from FlagEmbedding import FlagReranker
_model = None
_lock = False
MODEL_PATH = "/data/ymq/models/BAAI/bge-reranker-v2-m3"
DEVICE = "cuda:0"
def get_model():
global _model, _lock
if _model is None and not _lock:
_lock = True
try:
_model = FlagReranker(MODEL_PATH, use_fp16=True, device=DEVICE)
except Exception as e:
_lock = False
raise RuntimeError(f"Failed to load BGE-Reranker: {e}")
return _model
def rerank(query: str, documents: list, top_k: int = 5) -> dict:
"""Rerank documents by relevance to query."""
model = get_model()
# Create pairs: [[query, doc1], [query, doc2], ...]
pairs = [[query, doc] for doc in documents]
# Compute scores
scores = model.compute_score(pairs, normalize=True)
if isinstance(scores, float):
scores = [scores]
# Pair with documents and sort
ranked = list(zip(scores, documents))
ranked.sort(key=lambda x: -x[0])
# Return top_k
results = [
{"doc": doc, "score": round(score, 6)}
for score, doc in ranked[:top_k]
]
return {"ranked_docs": results, "total": len(documents), "returned": len(results)}
def health_check():
"""Check model status."""
model = get_model()
return {
"model": "BAAI/bge-reranker-v2-m3",
"device": DEVICE,
"loaded": model is not None,
"model_path": MODEL_PATH
}