使用Text Embeddings Inference运行Embeddings模型和Rerank模型

TEI简介

Text Embeddings Inference（TEI）是一个全面的工具包，旨在高效地部署和服务开源文本嵌入模型。它支持对最流行的模型进行高性能提取，包括FlagEmbedding、Ember、GTE和E5。

TEI提供多种功能，旨在优化部署过程并提高整体性能。

简化部署：TEI消除了模型图(model graph)编译步骤，简化了部署过程。
高效资源利用：得益于小的Docker镜像和快速的启动时间，可以实现真正的无服务器功能。
动态批处理：TEI采用基于token的动态批处理，从而在推理过程中优化资源利用。
优化的推理：TEI通过使用优化的Transformers代码来推理，利用Flash Attention、Candle和cuBLASLt。
Safetensors权重加载： TEI加载Safetensors权重以加快启动时间。
生产就绪：TEI支持通过Open Telemetry进行分布式跟踪，并导出Prometheus指标。
支持Metal，可以在Mac上本地执行

Metal

Metal是由苹果公司开发的一种高性能图形和计算API（应用程序编程接口），它可以直接访问Mac设备上的图形处理单元（GPU），并提供低级别、高效率的计算控制。Metal 主要用于图形渲染（如游戏和3D应用程序），但也可以用于通用计算任务，比如机器学习中的大规模并行计算。

在本地使用TEI

本地使用TEI需要从源码编译安装TEI。TEI是用Rust语言开发的，需要确保本地已经安装了rust。

cargo --version

可参考“安装Rust”

下载TEI的源码：

git clone https://github.com/huggingface/text-embeddings-inference.git
cd text-embeddings-inference
# 将使用tag v1.5.0构建
git checkout -b tobuild v1.5.0

在本地CPU使用TEI

操作系统Ubuntu 24.04, CPU架构x86_64。

sudo apt install pkg-config

sudo apt install intel-mkl

cargo install --path router -F mkl

cargo install --path router -F mkl

-F 是 cargo 命令的一个选项，用于启用某个特性（feature）。-F 是 --features 的缩写形式，意思是为构建过程启用特定的可选功能或配置。

mkl表示启用Intel Math Kernel Library (MKL)，这是一个高性能数学库，通常用于线性代数、FFT、稀疏矩阵运算等计算密集型任务。通过启用这个特性，会使用MKL来优化性能，尤其是在涉及到数值计算和机器学习推理时。

Intel Math Kernel Library (MKL)不是完全免费的，所以放弃“在本地CPU使用TEI”的方式

在本地GPU使用TEI

操作系统Ubuntu 24.04, CPU架构x86_64, NVIDIA GeForce RTX 4090。

确保您已安装CUDA和NVIDIA驱动程序，设备上的NVIDIA驱动程序需要与CUDA 12.2或更高版本兼容。将NVIDIA可执行文件添加到的环境变量路径中：

export PATH=$PATH:/usr/local/cuda/bin

对于Turing架构GPU（T4、RTX 2000系列…）

sudo apt install pkg-config

cargo install --path router -F candle-cuda-turing -F http --no-default-features

对于For Ampere and Hopper：

sudo apt install pkg-config

cargo install --path router -F candle-cuda -F http --no-default-features

这里是4090显卡，使用下面的命令编译安装：

cargo install --path router -F candle-cuda -F http --no-default-features

安装成功后，使用下面的命令查看TEI的版本号：

text-embeddings-router -V
text-embeddings-router 1.5.0

运行嵌入模型BAAI/bge-large-zh-v1.5

预先下载模型BAAI/bge-large-zh-v1.5：

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download BAAI/bge-large-zh-v1.5
...
~/.cache/huggingface/hub/models--BAAI--bge-large-zh-v1.5/snapshots/79e7739b6ab944e86d6171e44d24c997fc1e0116

运行嵌入模型：

export HF_HUB_OFFLINE=1
export HF_ENDPOINT=https://hf-mirror.com
text-embeddings-router --model-id BAAI/bge-large-zh-v1.5 --revision 79e7739b6ab944e86d6171e44d24c997fc1e0116 --port 8080

运行上面的命令，无论HF_HUB_OFFLINE设置成什么都会去下载1_Pooling/config.json，而且TEI好像不认HF_ENDPOINT这个环境变量。并且程序卡在下面的那行日志上：

INFO download_pool_config: text_embeddings_core::download: core/src/download.rs:38: Downloading `1_Pooling/config.json`

当前的解决方案是，直接设置--model-id为模型目录。

text-embeddings-router --model-id ~/.cache/huggingface/hub/models--BAAI--bge-large-zh-v1.5/snapshots/79e7739b6ab944e86d6171e44d24c997fc1e0116 \
  --port 8080

INFO text_embeddings_router: router/src/main.rs:175: Args { model_id: "/hom*/******/.*****/***********/***/******--****--***-*****-**-**.*/*********/*************************************116", revision: None, tokenization_workers: None, dtype: None, pooling: None, max_concurrent_requests: 512, max_batch_tokens: 16384, max_batch_requests: None, max_client_batch_size: 32, auto_truncate: false, default_prompt_name: None, default_prompt: None, hf_api_token: None, hostname: "0.0.0.0", port: 8080, uds_path: "/tmp/text-embeddings-inference-server", huggingface_hub_cache: None, payload_limit: 2000000, api_key: None, json_output: false, otlp_endpoint: None, otlp_service_name: "text-embeddings-inference.server", cors_allow_origin: None }
INFO text_embeddings_router: router/src/lib.rs:199: Maximum number of tokens per request: 512
INFO text_embeddings_core::tokenization: core/src/tokenization.rs:28: Starting 12 tokenization workers
INFO text_embeddings_router: router/src/lib.rs:241: Starting model backend
INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:251: Starting FlashBert model on Cuda(CudaDevice(DeviceId(1)))
INFO text_embeddings_router: router/src/lib.rs:257: Warming up model
INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:8080
INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

可以使用http://127.0.0.1:8080/docs查看接口的Swagger文档。

可以看到暴露了与OPENAI兼容的嵌入API端点/v1/embeddings，使用curl测试如下：

curl --location --request POST 'http://127.0.0.1:8080/v1/embeddings' \
-H 'Content-Type: application/json' \
-d '{
  "input": ["大模型是什么"]
}'

{"object":"list","data":[{"object":"embedding","embedding":[-0.010773913,0.02412172,-0.0015916008,0.044735655,-0.0074426555,...

运行Rerank模型BAAI/bge-reranker-v2-m3

预先下载模型BAAI/bge-reranker-v2-m3：

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download BAAI/bge-reranker-v2-m3
~/.cache/huggingface/hub/models--BAAI--bge-reranker-v2-m3/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e

运行Rerank模型：

export HF_HUB_OFFLINE=1
export HF_ENDPOINT=https://hf-mirror.com
text-embeddings-router --model-id BAAI/bge-reranker-v2-m3 --revision 953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e --port 8081

INFO download_pool_config: text_embeddings_core::download: core/src/download.rs:38: Downloading `1_Pooling/config.json`

当前的解决方案是，直接设置--model-id为模型目录。

text-embeddings-router --model-id ~/.cache/huggingface/hub/models--BAAI--bge-reranker-v2-m3/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e \
  --port 8081

INFO text_embeddings_router: router/src/main.rs:175: Args { model_id: "/hom*/******/.*****/***********/***/******--****--***-********-**-**/*********/*************************************41e", revision: None, tokenization_workers: None, dtype: None, pooling: None, max_concurrent_requests: 512, max_batch_tokens: 16384, max_batch_requests: None, max_client_batch_size: 32, auto_truncate: false, default_prompt_name: None, default_prompt: None, hf_api_token: None, hostname: "0.0.0.0", port: 8081, uds_path: "/tmp/text-embeddings-inference-server", huggingface_hub_cache: None, payload_limit: 2000000, api_key: None, json_output: false, otlp_endpoint: None, otlp_service_name: "text-embeddings-inference.server", cors_allow_origin: None }
WARN text_embeddings_router: router/src/lib.rs:195: Could not find a Sentence Transformers config
INFO text_embeddings_router: router/src/lib.rs:199: Maximum number of tokens per request: 8192
INFO text_embeddings_core::tokenization: core/src/tokenization.rs:28: Starting 12 tokenization workers
INFO text_embeddings_router: router/src/lib.rs:241: Starting model backend
INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:285: Starting FlashBert model on Cuda(CudaDevice(DeviceId(1)))
INFO text_embeddings_router: router/src/lib.rs:257: Warming up model
INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:8081
INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

Rerank API的端点是/rerank，使用curl测试如下：

curl -s -X POST http://127.0.0.1:8081/rerank \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
    "query": "中国的首都是哪里？",
    "return_text": true,
    "texts": [
    "中国的首都是北京，是国家的政治、文化和国际交往中心。",
    "中国有许多历史名城，其中包括北京、上海、广州等城市。",
    "北京是中国的首都，拥有丰富的历史文化遗产，如故宫和长城。",
    "日本的首都是东京，东京是日本最大的城市之一。",
    "北京位于中国的华北平原，是中国的政治中心。",
    "上海是中国的金融中心，位于中国的东部沿海。",
    "在中国的历史上，曾有多个城市作为首都，包括南京、洛阳、西安等。",
    "北京自明朝以来一直是中国的首都。",
    "中国的主要城市包括北京、上海、广州、深圳等。",
    "北京作为首都，在经济、文化、国际关系中占据重要位置。"
    ]
  }' | jq .

输出内容如下：

[
  {
    "index": 0,
    "text": "中国的首都是北京，是国家的政治、文化和国际交往中心。",
    "score": 0.9986853
  },
  {
    "index": 7,
    "text": "北京自明朝以来一直是中国的首都。",
    "score": 0.9857729
  },
  {
    "index": 2,
    "text": "北京是中国的首都，拥有丰富的历史文化遗产，如故宫和长城。",
    "score": 0.95806944
  },
  {
    "index": 8,
    "text": "中国的主要城市包括北京、上海、广州、深圳等。",
    "score": 0.95719784
  },
  {
    "index": 6,
    "text": "在中国的历史上，曾有多个城市作为首都，包括南京、洛阳、西安等。",
    "score": 0.95695716
  },
  {
    "index": 4,
    "text": "北京位于中国的华北平原，是中国的政治中心。",
    "score": 0.77149934
  },
  {
    "index": 9,
    "text": "北京作为首都，在经济、文化、国际关系中占据重要位置。",
    "score": 0.6960617
  },
  {
    "index": 1,
    "text": "中国有许多历史名城，其中包括北京、上海、广州等城市。",
    "score": 0.642881
  },
  {
    "index": 5,
    "text": "上海是中国的金融中心，位于中国的东部沿海。",
    "score": 0.33709443
  },
  {
    "index": 3,
    "text": "日本的首都是东京，东京是日本最大的城市之一。",
    "score": 0.022032265
  }
]

命令行参数

text-embeddings-router命令行参数，详细见文档。

容器化运行TEI

model=BAAI/bge-large-zh-v1.5
volume=$HOME/.cache/huggingface
# sudo docker run --name tei-embedding --restart=unless-stopped -d --gpus all --shm-size 2g -p 8080:80 -v $volume:/data \
#     -e HF_HUB_OFFLINE=1 \
#     -e HF_ENDPOINT=https://hf-mirror.com \
#     ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
#     --model-id $model


sudo docker run --name tei-embedding --restart=unless-stopped -d --gpus all --shm-size 2g -p 8080:80 -v $volume:/data \
    -e HF_HUB_OFFLINE=1 \
    -e HF_ENDPOINT=https://hf-mirror.com \
    ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
    --model-id /data/hub/models--BAAI--bge-large-zh-v1.5/snapshots/79e7739b6ab944e86d6171e44d24c997fc1e0116

model=BAAI/bge-reranker-v2-m3
volume=$HOME/.cache/huggingface
# sudo docker run --name tei-reranker --restart=unless-stopped -d --gpus all --shm-size 2g -p 8080:80 -v $volume:/data \
#     -e HF_HUB_OFFLINE=1 \
#     -e HF_ENDPOINT=https://hf-mirror.com \
#     ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
#     --model-id $model


sudo docker run --name tei-reranker --restart=unless-stopped -d --gpus all --shm-size 2g -p 8081:80 -v $volume:/data \
    -e HF_HUB_OFFLINE=1 \
    -e HF_ENDPOINT=https://hf-mirror.com \
    ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
    --model-id /data/hub/models--BAAI--bge-reranker-v2-m3/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e