From a4791eb1e00395768a6796ee943ba1688e7cba45 Mon Sep 17 00:00:00 2001 From: yumoqing Date: Mon, 15 Jun 2026 08:32:08 +0800 Subject: [PATCH] feat: BGE-Reranker service with /api/status and /api/rerank endpoints --- .gitignore | 7 ++++ README.md | 79 +++++++++++++++++++++++++++++++++++++++++ ah.py | 6 ++++ build.sh | 57 +++++++++++++++++++++++++++++ conf/config.json | 25 +++++++++++++ init.py | 55 ++++++++++++++++++++++++++++ workers/__init__.py | 0 workers/bge_reranker.py | 60 +++++++++++++++++++++++++++++++ 8 files changed, 289 insertions(+) create mode 100644 .gitignore create mode 100644 README.md create mode 100644 ah.py create mode 100755 build.sh create mode 100644 conf/config.json create mode 100644 init.py create mode 100644 workers/__init__.py create mode 100644 workers/bge_reranker.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1ec227d --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +ah.pid +nohup.out +__pycache__/ +*.pyc +workers/__pycache__/ +logs/ +files/ diff --git a/README.md b/README.md new file mode 100644 index 0000000..5b6b6d3 --- /dev/null +++ b/README.md @@ -0,0 +1,79 @@ +# 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 diff --git a/ah.py b/ah.py new file mode 100644 index 0000000..fd4f109 --- /dev/null +++ b/ah.py @@ -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) diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..7aee39f --- /dev/null +++ b/build.sh @@ -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 diff --git a/conf/config.json b/conf/config.json new file mode 100644 index 0000000..7a6f176 --- /dev/null +++ b/conf/config.json @@ -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"] + ] + } +} diff --git a/init.py b/init.py new file mode 100644 index 0000000..4beb541 --- /dev/null +++ b/init.py @@ -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) diff --git a/workers/__init__.py b/workers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/workers/bge_reranker.py b/workers/bge_reranker.py new file mode 100644 index 0000000..9106e6d --- /dev/null +++ b/workers/bge_reranker.py @@ -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 + }