本地Rerank服务实现方案

在构建RAG系统时，需要为Dify配置本地的重排序（Rerank）服务。目标：部署一个稳定、支持中文优化的rerank服务，供Dify调用。

🔄 技术路线演进

第一阶段：LocalAI Docker方案（尝试后放弃）

初始尝试

# 错误尝试：镜像标签问题
docker run -p 8035:8080 -v ./models:/models --name localai quay.io/go-skynet/local-ai:latest-ffmpeg
# 错误：manifest unknown

# 正确镜像（Mac M1）
docker run -d -p 8035:8080 \
  -v $(pwd)/models:/models \
  -v $(pwd)/config:/config \
  --name localai \
  localai/localai:latest-aio-cpu

遇到的问题

1. 后端缺失错误：backend not found: bert-embeddings
2. 模型加载失败：failed to load model with internal loader
3. 自动下载依赖：下载语音模型耗时极长（显示15小时）

部分成功发现

• jina-reranker-v1-base-en 模型可以正常工作
• 但该模型对中文优化不足

# 验证jina模型可用
curl -X POST http://localhost:8035/v1/rerank \
  -H "Content-Type: application/json" \
  -d '{
    "model": "jina-reranker-v1-base-en",
    "query": "What is AI?",
    "documents": ["doc1", "doc2"],
    "top_n": 3
  }'
# 返回成功，但中文效果一般

第二阶段：Python原生服务（最终采用）

核心实现代码

# rerank_svr.py
import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # 关键：国内镜像

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import CrossEncoder
import uvicorn
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

# 模型配置
MODEL_MAP = {
    "bge-chinese": "BAAI/bge-reranker-v2-m3",  # 中文优化版本
    "bge-base": "BAAI/bge-reranker-base",
    "minilm": "cross-encoder/ms-marco-MiniLM-L-6-v2"
}

models = {}

def load_model(model_name="bge-chinese"):
    """加载模型，处理网络问题"""
    if model_name in models:
        return models[model_name]
    
    logger.info(f"加载模型: {model_name}")
    try:
        model_path = MODEL_MAP[model_name]
        models[model_name] = CrossEncoder(model_path)
        logger.info(f"模型加载成功: {model_name}")
        return models[model_name]
    except Exception as e:
        logger.error(f"模型加载失败: {e}")
        raise HTTPException(status_code=500, detail=f"模型加载失败: {e}")

class RerankRequest(BaseModel):
    model: str = "bge-chinese"
    query: str
    documents: list[str]
    top_n: int = None

@app.post("/v1/rerank")
async def rerank(request: RerankRequest):
    """重排序API端点"""
    try:
        model = load_model(request.model)
        
        # 准备查询-文档对
        pairs = [(request.query, doc) for doc in request.documents]
        
        # 计算相关性分数
        scores = model.predict(pairs).tolist()
        
        # 格式化结果
        results = []
        for i, (doc, score) in enumerate(zip(request.documents, scores)):
            results.append({
                "index": i,
                "document": {"text": doc},
                "relevance_score": float(score)
            })
        
        # 排序
        results.sort(key=lambda x: x["relevance_score"], reverse=True)
        
        # 应用top_n
        if request.top_n:
            results = results[:request.top_n]
        
        return {
            "model": request.model,
            "results": results,
            "usage": {"total_tokens": len(request.query) + sum(len(d) for d in request.documents)}
        }
        
    except Exception as e:
        logger.error(f"重排序失败: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/readyz")
async def health():
    return {"status": "ready"}

@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {"id": "bge-chinese", "object": "model"},
            {"id": "bge-base", "object": "model"},
            {"id": "minilm", "object": "model"}
        ]
    }

if __name__ == "__main__":
    logger.info("启动BGE中文重排序服务...")
    uvicorn.run(app, host="0.0.0.0", port=8036)

部署与运行

# 安装依赖
pip install fastapi uvicorn sentence-transformers pydantic

# 运行服务
python rerank_svr_fixed.py

# 测试服务
curl http://localhost:8036/readyz
curl http://localhost:8036/v1/models

# 测试中文重排序
curl -X POST http://localhost:8036/v1/rerank \
  -H "Content-Type: application/json" \
  -d '{
    "model": "bge-chinese",
    "query": "什么是人工智能？",
    "documents": [
      "人工智能是计算机科学的分支",
      "机器学习通过数据训练模型",
      "Python是编程语言"
    ],
    "top_n": 2
  }'

🔑 关键问题与解决方案

1. 网络问题：HuggingFace连接失败

问题现象：

SSLError(SSLEOFError(8, '[SSL: UNEXPECTED_EOF_WHILE_READING] EOF occurred in violation of protocol')

解决方案：

import os
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'  # 使用国内镜像

2. Docker网络连通性问题

问题现象：Dify容器无法访问宿主机服务

Connection refused (Errno 111)

解决方案：

• 在Dify中使用：http://host.docker.internal:8036/v1/rerank
• 不是：http://localhost:8036/v1/rerank（容器内指向自己）
• 不是：http://127.0.0.1:8036/v1/rerank

3. 服务管理脚本

#!/bin/bash
# 启动服务
export HF_ENDPOINT=https://hf-mirror.com
nohup python rerank_svr_fixed.py > rerank.log 2>&1 &

# 检查服务
curl http://localhost:8036/readyz

# 查看日志
tail -f rerank.log

📋 在Dify中的配置

1. 进入Dify设置：工具 → 外部API工具
2. 配置参数：
   • API地址：http://host.docker.internal:8036/v1/rerank
   • 模型名称：bge-chinese
   • API密钥：留空
3. 测试连接：使用简单查询测试

🎯 模型对比实际测试结果

测试案例：中文查询

# 测试查询："什么是机器学习？"
# 测试文档：
# 1. "人工智能的一个子领域"
# 2. "通过数据训练模型" 
# 3. "深度学习使用神经网络"
# 4. "Python编程语言"

# BGE中文模型结果（更合理）：
# 1. "通过数据训练模型" (分数最高)
# 2. "人工智能的一个子领域"
# 3. "深度学习使用神经网络"

# Jina模型结果：
# 1. "深度学习使用神经网络" 
# 2. "通过数据训练模型"
# 3. "人工智能的一个子领域"

实际使用建议

1. 中文场景：使用 bge-chinese (BAAI/bge-reranker-v2-m3)
2. 英文场景：使用 minilm (cross-encoder/ms-marco-MiniLM-L-6-v2)
3. 混合场景：使用 bge-base (BAAI/bge-reranker-base)

⚠️ 注意事项

1. 首次运行：需要下载模型文件，国内网络需配置镜像
2. 内存占用：BGE模型约占用1-2GB内存
3. 端口冲突：确保8036端口未被占用
4. 依赖版本：建议使用较新的sentence-transformers版本

✅ 验证服务正常

服务正常运行时应有以下特征：

1. 启动日志显示模型加载成功
2. /readyz 端点返回 {"status": "ready"}
3. /v1/rerank 能返回合理的相关性分数
4. 在Dify中测试连接成功

📝 总结

经过实践验证，Python原生服务方案相比LocalAI Docker方案更稳定可靠，主要原因：

1. 依赖简单：仅需Python环境，无需复杂Docker配置
2. 调试方便：错误信息清晰，容易排查
3. 网络可控：可灵活配置镜像源
4. 资源占用少：无额外的容器开销

当前方案已稳定运行，为Dify提供中文优化的重排序能力。后续可根据实际使用情况考虑性能优化和监控增强。

---

创建时间：2025-12-17 08:28