下载模型#

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 的核心任务是把正确的排在前面，把错误的踢到后面。它完美地把笔记本踢到了最后一名。

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

下载模型#

配置 onnx 转换#

配置 docker-compose#

测试#

评论