搭建TEI平台并部署模型

BennuCTech
673 阅读5分钟

前言

TEI就是text-embeddings-inference，github地址是github.com/huggingface…

TEI 是一个用于部署和提供开源文本嵌入和序列分类模型的工具，该工具主要是以部署 Embedding 模型为主，但是也支持 Rerank 和其他类型的模型的部署，同时它还支持部署兼容 OpenAI API 的 API 服务。

TEI是HuggingFace推出的，官网地址是huggingface.co/ 。

Docker部署

TEI支持Docker直接部署模型，命令如下：

docker run --gpus all -p 8080:80 -v $volume:/data --pull always ghcr.io/huggingface/text-embeddings-inference:1.5 --model-id $model

这里$volume是本机上的一个地址，是给docker容器用于存储下载的文件，比如可以是~/data。

$model就是要部署的模型ID，比如我们打算部署一个rerank模型，就可以是BAAI/bge-reranker-large

注意命令中的text-embeddings-inference:1.5是TEI的Image，但是它提供了几个Image，分别对应不同的GPU系列，所以要选择自己适合的Image。

模型

TEI支持的模型如下

Text Embeddings
MTEB RankModel SizeModel TypeModel ID
17B (Very Expensive)MistralSalesforce/SFR-Embedding-2_R
27B (Very Expensive)Qwen2Alibaba-NLP/gte-Qwen2-7B-instruct
91.5B (Expensive)Qwen2Alibaba-NLP/gte-Qwen2-1.5B-instruct
150.4BAlibaba GTEAlibaba-NLP/gte-large-en-v1.5
200.3BBertWhereIsAI/UAE-Large-V1
240.5BXLM-RoBERTaintfloat/multilingual-e5-large-instruct
N/A0.1BNomicBertnomic-ai/nomic-embed-text-v1
N/A0.1BNomicBertnomic-ai/nomic-embed-text-v1.5
N/A0.1BJinaBERTjinaai/jina-embeddings-v2-base-en
N/A0.1BJinaBERTjinaai/jina-embeddings-v2-base-code
Rerank
TaskModel TypeModel ID
Re-RankingXLM-RoBERTaBAAI/bge-reranker-large
Re-RankingXLM-RoBERTaBAAI/bge-reranker-base
Sentiment AnalysisRoBERTaSamLowe/roberta-base-go_emotions

Docker Image

TEI提供的Image如下：

ArchitectureImage
CPUghcr.io/huggingface/text-embeddings-inference:cpu-1.5
VoltaNOT SUPPORTED
Turing (T4, RTX 2000 series, ...)ghcr.io/huggingface/text-embeddings-inference:turing-1.5 (experimental)
Ampere 80 (A100, A30)ghcr.io/huggingface/text-embeddings-inference:1.5
Ampere 86 (A10, A40, ...)ghcr.io/huggingface/text-embeddings-inference:86-1.5
Ada Lovelace (RTX 4000 series, ...)ghcr.io/huggingface/text-embeddings-inference:89-1.5
Hopper (H100)ghcr.io/huggingface/text-embeddings-inference:hopper-1.5 (experimental)

对应自己的GPU系列来选择合适的Image，如果GPU太低就选择CPU版本的即可。

实例

以测试电脑为例，GPU太低所以只能用text-embeddings-inference:cpu-1.5，所以部署BAAI/bge-reranker-large这个模型的命令是

docker run -p 8080:80 -v ～/data:/data --pull always ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 --model-id BAAI/bge-reranker-large

注意因为我们不使用GPU，所以命令中的--gpus all可以去掉，如果使用就需要根据自己的需要添加上。

部署中需要下载，最好有个代理，以免下载失败。

部署需要一定时间，完成后会看到如下信息：

INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:80
INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

这时候才是部署完成，可以正常使用。

模型下载失败

在部署的过程中发现报错

Error: Could not create backend

Caused by:
    Could not start backend: Failed to create ONNX Runtime session: Deserialize tensor roberta.encoder.layer.20.attention.self.value.bias failed.GetFileLength for /data/models--BAAI--bge-reranker-large/snapshots/55611d7bca2a7133960a6d3b71e083071bbfc312/onnx/model.onnx_data failed:Invalid fd was supplied: -1

向上查看详细日志可以知道是模型的model.onnx和model.onnx_data这两个文件下载失败，日志中看地址是huggingface.co/BAAI/bge-re…

访问这个地址已经404了，其他模型也是同样的问题，不知道是不是官方的Image太旧了而且地址换了。

我在官网上手动下载了这两个文件，在huggingface.co/models 搜索该模型并点进去，然后点击Files and versions，在下面目录中进入onnx，地址是huggingface.co/BAAI/bge-re… 。这样就可以看到这两个文件了，手动下载到本地。

下载完成后将这两个文件拷贝到$volume/models--BAAI--bge-reranker-large/snapshots/55611d7bca2a7133960a6d3b71e083071bbfc312/onnx/下，这里$volume就是我们在命令里配置的地址，models--BAAI--bge-reranker-large就是模型名称，进去看一下就知道了。55611d7bca2a7133960a6d3b71e083071bbfc312是快照名称，同样进去目录看一下就知道了。

然后重新执行命令就可以了。

后台运行

上面的命令会阻塞终端，停止终端docker容器也会停止。如果想后台运行容器，可以添加-d，如下：

docker run -d -p 8080:80 --name TEI -v ～/data:/data --pull always ghcr.io/huggingface/text-embeddings-inference:cpu-1.5 --model-id BAAI/bge-reranker-large

我们这里同时也为容器命名为TEI，如果不命名就会随机名字，有了这个名字我们后续停止重启都方便。

停止容器docker stop TEI

重启容器docker start TEI

注意在后台运行就无法看到运行日志，这样是否完成或者有问题也不容易看出来，如果是windows上使用docker desktop可以直接点击容器看到日志。也可以使用docker logs TEI命令来查看某个容器的日志

使用

部署成功后可以通过http://localhost:8080/docs 来查看服务的API文档

image.png

测试

比如我们部署的是rerank模型，可以展开/rerank，这里可以进行测试

7F6AE741-918E-4E34-9ABD-C6E53AE480E3.png

点击 try it out，下面就会自动生成一个请求内容的示例，如下

AEA11C93-B4AE-4419-A010-AC76589AF981.png

在点击 Execute 就会发送这个测试请求，然后下面可以看到返回结果和curl等内容。

注意：rerank这里的示例有问题，最后一个truncation_direction的值应该是Left或者Right，不能小写，所以执行示例的话会报错

Failed to deserialize the JSON body into the target type: truncation_direction: unknown variant `right`, expected `Left` or `Right` at line 9 column 33

改正确即可。

请求

测试通过我们就使用这个接口了，比如rerank接口地址http://localhost:8080/rerank

然后需要在header中添加

accept: application/json
Content-Type: application/json

请求body如下：

{
  "query": "What is Deep Learning?",
  "raw_scores": false,
  "return_text": false,
  "texts": [
    "Deep Learning is ..."
  ],
  "truncate": false,
  "truncation_direction": "Right"
}

得到的结果如下：

[    {        "index": 0,        "score": 0.9972958    }]

对比query给texts中的列表重新评分，然后通过这个score就可以重新进行排序了。

avatar
文章
阅读
粉丝