Jacin Blog
AI 大模型开发

tei-rerank 部署重排模型bge-reranker-v2-m3

JACIN7 分钟阅读
AI 摘要

这份记录聚焦在用 Hugging Face Text Embeddings Inference 的 CPU 镜像部署 BAAI/bge-reranker-v2-m3 重排模型，场景是为检索结果增加一个独立的 rerank 服务，而不是普通 embedding 服务。流程先通过 snapshot_download 将模型完整下载到本地目录，再安装 optimum、onnx、onnxruntime，并用 ORTModelForSequenceClassification.from_pretrained(export=True) 将 PyTorch 模型导出到模型目录下的 onnx 子目录，同时保存 tokenizer。docker-compose 配置的关键点包括使用 ghcr.io/huggingface/text-embeddings-inference:cpu-1.5，挂载本地模型目录到 /data/rerank，开放 38190:80 端口，并通过 --model-id 指向容器内路径，配合 --auto-truncate 和 --dtype float32 启动服务。文章特别强调 shm_size: 1g 的必要性，因为 Docker 默认共享内存只有 64MB，AI 推理程序在并发或大矩阵数据传输时可能因共享内存不足直接 Bus error 崩溃。最后用 /rerank 接口测试“我想买一个手机”与三星、苹果、联想笔记本三段文本的相关性排序，结果显示手机相关文本得分约 0.3，而笔记本得分接近 0.0001。读者可以从中获得一套可复现的 TEI rerank 部署步骤，并理解重排模型更应关注候选项之间的相对差距和排序效果，而不是孤立解读绝对分值。

下载模型#

text 
import os

save_dir = "/root/tei/data/bge-reranker-v2-m3" # 存到新目录
os.makedirs(save_dir, exist_ok=True)

print(f"🚀 正在下载 Rerank 模型到 {save_dir} ...")

snapshot_download(
    repo_id="BAAI/bge-reranker-v2-m3",
    local_dir=save_dir,
    local_dir_use_symlinks=False,
    resume_download=True
)
print("✅ 下载完成！")

配置 onnx 转换#

text 
pip install optimum[exporters] onnx onnxruntime
pip install --upgrade "optimum[onnxruntime]"

vim convert_to_onnx.py 下载模型，因为需要 onnx 转换，所以使用 python 进行转换

text 
import os
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

# 1. 设置路径
input_model_path = "./data/bge-reranker-v2-m3"  # 你的 PyTorch 模型路径
output_onnx_path = "./data/bge-reranker-v2-m3/onnx" # 输出的 ONNX 路径

print(f"🚀 正在加载模型: {input_model_path} ...")
print("⏳ 正在进行 ONNX 转换，这可能需要几分钟，请耐心等待...")

try:
    # 2. 加载并转换 (export=True 是关键)
    model = ORTModelForSequenceClassification.from_pretrained(
        input_model_path,
        export=True
    )
    tokenizer = AutoTokenizer.from_pretrained(input_model_path)

    # 3. 保存到指定目录
    print(f"💾 正在保存 ONNX 模型到: {output_onnx_path} ...")
    model.save_pretrained(output_onnx_path)
    tokenizer.save_pretrained(output_onnx_path)

    print("✅ 转换成功！你现在可以启动 Docker 了！")

except Exception as e:
    print(f"❌ 转换失败: {e}")

配置 docker-compose#

text 
  # 服务 2: 重排序 (Rerank) - 新增的！
  tei-rerank:
    container_name: tei-rerank
    image: ghcr.io/huggingface/text-embeddings-inference:cpu-1.5
    restart: always
    ports:
      - "38190:80"  # ⚠️ 端口不一样，别冲突
    volumes:
      # 👇 挂载刚才下载的 rerank 模型目录
      - ./data/bge-reranker-v2-m3:/data/rerank
    shm_size: 1g
    environment:
      - no_proxy=*
    command: 
      - --model-id
      - /data/rerank  # 指向容器内的挂载路径
        # 👇 必须加上这个！告诉 TEI 这是一个 Rerank 模型，不是 Embedding
      - --auto-truncate
      # 👇 关键：如果不加这个，某些版本的 TEI 会死命找 ONNX
      - --dtype
      - float32
  • Docker 的坑： 默认情况下，Docker 只给容器分配 64MB 的共享内存。
  • AI 的需求： TEI、PyTorch、Ollama 这些 AI 程序，在处理大量并发或大模型时，需要在进程间传输巨大的矩阵数据。
  • 后果： 如果不加这行代码，共享内存只有 64MB，一旦不够用，容器就会直接崩溃报错（通常报 Bus error），而不是慢一点那么简单。

测试#

text 
curl --request POST \
  --url http://127.0.0.1:38190/rerank \
  --header 'content-type: application/json' \
  --data '{
  "query": "我想买一个手机",
  "texts": [
    "三星 Galaxy S24 Ultra 是一款顶级的安卓智能，拍照效果极佳。",
    "苹果 iPhone 15 Pro Max ，搭载 A17 芯片，性能强劲。",
    "联想拯救者笔记本电脑，搭载 RTX4060 显卡。"
  ]
}
'

结果：

text 
[
  {
    "index": 1,
    "score": 0.3314397
  },
  {
    "index": 0,
    "score": 0.30318388
  },
  {
    "index": 2,
    "score": 0.00010000849
  }
]

在 Rerank 领域，我们要看的是**“相对差距” (Relative Gap)**

  • 手机 (三星/苹果): ~0.3
  • 笔记本 (联想): 0.00008

算一下倍数：0.3 / 0.00008 = 3750 倍。 含义： 模型认为“三星 9800”是手机的概率，比“联想笔记本”是手机的概率，高出 3750 倍。 这就足够了！Rerank 的核心任务是把正确的排在前面，把错误的踢到后面。它完美地把笔记本踢到了最后一名。

评论

还没有评论，来发第一个吧