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

2024-10-21

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。

1cargo --version

可参考“安装Rust”

下载TEI的源码：

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

在本地CPU使用TEI #

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

1sudo apt install pkg-config
2
3sudo apt install intel-mkl
4
5cargo 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可执行文件添加到的环境变量路径中：

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

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

1sudo apt install pkg-config
2
3cargo install --path router -F candle-cuda-turing -F http --no-default-features

对于For Ampere and Hopper：

1sudo apt install pkg-config
2
3cargo install --path router -F candle-cuda -F http --no-default-features

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

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

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

1text-embeddings-router -V
2text-embeddings-router 1.5.0

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

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

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

运行嵌入模型：

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

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

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

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

 1text-embeddings-router --model-id ~/.cache/huggingface/hub/models--BAAI--bge-large-zh-v1.5/snapshots/79e7739b6ab944e86d6171e44d24c997fc1e0116 \
 2  --port 8080
 3
 4INFO 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 }
 5INFO text_embeddings_router: router/src/lib.rs:199: Maximum number of tokens per request: 512
 6INFO text_embeddings_core::tokenization: core/src/tokenization.rs:28: Starting 12 tokenization workers
 7INFO text_embeddings_router: router/src/lib.rs:241: Starting model backend
 8INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:251: Starting FlashBert model on Cuda(CudaDevice(DeviceId(1)))
 9INFO text_embeddings_router: router/src/lib.rs:257: Warming up model
10INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:8080
11INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

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

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

1curl --location --request POST 'http://127.0.0.1:8080/v1/embeddings' \
2-H 'Content-Type: application/json' \
3-d '{
4  "input": ["大模型是什么"]
5}'
6
7{"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：

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

运行Rerank模型：

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

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

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

 1text-embeddings-router --model-id ~/.cache/huggingface/hub/models--BAAI--bge-reranker-v2-m3/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e \
 2  --port 8081
 3
 4INFO 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 }
 5WARN text_embeddings_router: router/src/lib.rs:195: Could not find a Sentence Transformers config
 6INFO text_embeddings_router: router/src/lib.rs:199: Maximum number of tokens per request: 8192
 7INFO text_embeddings_core::tokenization: core/src/tokenization.rs:28: Starting 12 tokenization workers
 8INFO text_embeddings_router: router/src/lib.rs:241: Starting model backend
 9INFO text_embeddings_backend_candle: backends/candle/src/lib.rs:285: Starting FlashBert model on Cuda(CudaDevice(DeviceId(1)))
10INFO text_embeddings_router: router/src/lib.rs:257: Warming up model
11INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:8081
12INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

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

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

输出内容如下：

 1[
 2  {
 3    "index": 0,
 4    "text": "中国的首都是北京，是国家的政治、文化和国际交往中心。",
 5    "score": 0.9986853
 6  },
 7  {
 8    "index": 7,
 9    "text": "北京自明朝以来一直是中国的首都。",
10    "score": 0.9857729
11  },
12  {
13    "index": 2,
14    "text": "北京是中国的首都，拥有丰富的历史文化遗产，如故宫和长城。",
15    "score": 0.95806944
16  },
17  {
18    "index": 8,
19    "text": "中国的主要城市包括北京、上海、广州、深圳等。",
20    "score": 0.95719784
21  },
22  {
23    "index": 6,
24    "text": "在中国的历史上，曾有多个城市作为首都，包括南京、洛阳、西安等。",
25    "score": 0.95695716
26  },
27  {
28    "index": 4,
29    "text": "北京位于中国的华北平原，是中国的政治中心。",
30    "score": 0.77149934
31  },
32  {
33    "index": 9,
34    "text": "北京作为首都，在经济、文化、国际关系中占据重要位置。",
35    "score": 0.6960617
36  },
37  {
38    "index": 1,
39    "text": "中国有许多历史名城，其中包括北京、上海、广州等城市。",
40    "score": 0.642881
41  },
42  {
43    "index": 5,
44    "text": "上海是中国的金融中心，位于中国的东部沿海。",
45    "score": 0.33709443
46  },
47  {
48    "index": 3,
49    "text": "日本的首都是东京，东京是日本最大的城市之一。",
50    "score": 0.022032265
51  }
52]

命令行参数 #

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

容器化运行TEI #

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

 1model=BAAI/bge-reranker-v2-m3
 2volume=$HOME/.cache/huggingface
 3# sudo docker run --name tei-reranker --restart=unless-stopped -d --gpus all --shm-size 2g -p 8080:80 -v $volume:/data \
 4#     -e HF_HUB_OFFLINE=1 \
 5#     -e HF_ENDPOINT=https://hf-mirror.com \
 6#     ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
 7#     --model-id $model
 8
 9
10sudo docker run --name tei-reranker --restart=unless-stopped -d --gpus all --shm-size 2g -p 8081:80 -v $volume:/data \
11    -e HF_HUB_OFFLINE=1 \
12    -e HF_ENDPOINT=https://hf-mirror.com \
13    ghcr.io/huggingface/text-embeddings-inference:1.5.1 \
14    --model-id /data/hub/models--BAAI--bge-reranker-v2-m3/snapshots/953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e