merge: bge-reranker service from GPU server
This commit is contained in:
commit
75e5bc85e6
7
.gitignore
vendored
Normal file
7
.gitignore
vendored
Normal file
@ -0,0 +1,7 @@
|
||||
ah.pid
|
||||
nohup.out
|
||||
__pycache__/
|
||||
*.pyc
|
||||
workers/__pycache__/
|
||||
logs/
|
||||
files/
|
||||
79
README.md
79
README.md
@ -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
6
ah.py
Normal 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
57
build.sh
Executable 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
25
conf/config.json
Normal 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
55
init.py
Normal 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
0
workers/__init__.py
Normal file
60
workers/bge_reranker.py
Normal file
60
workers/bge_reranker.py
Normal 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
|
||||
}
|
||||
Loading…
x
Reference in New Issue
Block a user