ghcr.io/huggingface/text-embeddings-inference:rocm-sha-5696458

Text Embeddings Inference

一个用于文本嵌入模型的极速推理解决方案。

在 NVIDIA A10 上使用序列长度为 512 个 token 的 https://huggingface.co/BAAI/bge-base-en-v1.5 基准测试：

目录

• Get Started
• Supported Models
• Docker
• Docker Images
• API Documentation
• Using a private or gated model
• Air gapped deployment
• Using Re-rankers models
• Using Sequence Classification models
• Using SPLADE pooling
• Distributed Tracing
• gRPC
• Local Install
• Apple Silicon (Homebrew)
• Docker Build
• ARM64 / aarch64
• AMD Instinct GPUs (ROCm)
• Examples

Text Embeddings Inference (TEI) 是一个用于部署和服务开源文本嵌入及序列分类模型的工具包。TEI 支持对最流行的模型进行高性能提取，包括 FlagEmbedding、Ember、GTE 和 E5。TEI 实现了许多功能，例如：

• 无需模型图编译步骤
• 支持 Metal，可在 Mac 上本地执行
• 体积小巧的 Docker 镜像和快速启动时间，为真正的无服务器架构做好准备！
• 基于 token 的动态批处理
• 使用 https://github.com/HazyResearch/flash-attention、https://github.com/huggingface/candle 和 https://docs.nvidia.com/cuda/cublas/#using-the-cublaslt-api 优化的推理用 transformers 代码
• https://github.com/huggingface/safetensors 权重加载
• https://github.com/onnx/onnx 权重加载
• 生产环境就绪（支持 Open Telemetry 分布式追踪、Prometheus 指标）

Get Started

Supported Models

Text Embeddings

Text Embeddings Inference 目前支持具有绝对位置编码的 Nomic、BERT、CamemBERT、XLM-RoBERTa 模型，具有 Alibi 位置编码的 JinaBERT 模型，具有 Rope 位置编码的 Mistral、Alibaba GTE、Qwen2 模型，以及 MPNet、ModernBERT、Qwen3 和 Gemma3。

以下是当前支持的部分模型示例：

要查看性能最佳的文本嵌入模型列表，请访问 https://huggingface.co/spaces/mteb/leaderboard。

Sequence Classification and Re-Ranking

Text Embeddings Inference 目前支持具有绝对位置编码的 CamemBERT 和 XLM-RoBERTa 序列分类模型。

以下是当前支持的部分模型示例：

Docker

model=Qwen/Qwen3-Embedding-0.6B
volume=$PWD/data # 与 Docker 容器共享卷以避免每次运行都下载权重

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

然后您可以发起如下请求：

curl 127.0.0.1:8080/embed \
-X POST \
-d '{"inputs":"What is Deep Learning?"}' \
-H 'Content-Type: application/json'

[!NOTE] 要使用 GPU，您需要安装 https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html。您机器上的 NVIDIA 驱动程序需要与 CUDA 12.2 或更高版本兼容。

要查看服务模型的所有选项：

$ text-embeddings-router --help
Text Embedding Webserver

Usage: text-embeddings-router [OPTIONS] --model-id

Options:
--model-id
Hugging Face 模型 ID，可以是任何带有 `text-embeddings-inference` 标签的模型（意味着它与 Text Embeddings Inference 兼容）。

或者，指定的ID也可以是包含由Transformers或Sentence Transformers的save_pretrained(...)方法保存的必要模型文件的本地目录路径。

[env: MODEL_ID=]

--revision
如果引用的是Hub上的模型，则为模型的实际修订版本。您可以使用特定的提交ID或分支，例如refs/pr/2

[env: REVISION=]

--tokenization-workers
可选地控制用于 payload 分词、验证和截断的分词器工作进程数。默认为机器上的CPU核心数

[env: TOKENIZATION_WORKERS=]

--dtype
要强制应用于模型的数据类型

[env: DTYPE=]
[possible values: float16, float32]

--served-model-name
正在提供服务的模型名称。如果未指定，默认为--model-id。仅用于通过HTTP的OpenAI兼容端点

[env: SERVED_MODEL_NAME=]

--pooling
可选地控制嵌入模型的池化方法。

如果未设置pooling，将从模型的1_Pooling/config.json配置中解析池化配置。

如果设置了pooling，它将覆盖模型的池化配置

[env: POOLING=]

可能的值：

• cls：选择CLS token作为嵌入
• mean：对模型嵌入应用均值池化
• splade：对模型嵌入应用SPLADE（稀疏词汇与扩展）。此选项仅在加载的模型是ForMaskedLM Transformer模型时可用
• last-token：选择最后一个token作为嵌入

--max-concurrent-requests
此特定部署的最大并发请求数。设置较低的限制将拒绝客户端请求，而不是让它们等待太久，通常有助于正确处理背压

[env: MAX_CONCURRENT_REQUESTS=]
[default: 512]

--max-batch-tokens

[!IMPORTANT] 这是允许最大程度利用可用硬件的关键控制项。

这表示一个批次内的潜在token总数。

对于max_batch_tokens=1000，您可以容纳10个total_tokens=100的查询，或一个1000个token的查询。

总体而言，这个数字应该尽可能大，直到模型受计算限制。由于实际内存开销取决于模型实现，text-embeddings-inference无法自动推断此数字。

[env: MAX_BATCH_TOKENS=]
[default: ***]

--max-batch-requests
可选地控制一个批次中的最大独立请求数

[env: MAX_BATCH_REQUESTS=]

--max-client-batch-size
控制客户端在单个请求中可以发送的最大输入数量

[env: MAX_CLIENT_BATCH_SIZE=]
[default: 32]

--auto-truncate
控制对超过模型最大支持大小的输入进行自动截断。默认为true（启用截断）。设置为false以禁用截断；禁用时，如果模型的最大输入长度超过--max-batch-tokens，服务器将拒绝启动并报错，而不是静默截断序列。

gRPC服务器不使用此参数

[env: AUTO_TRUNCATE=]

--default-prompt-name
默认用于编码的提示名称。如果未设置，则不应用任何提示。

必须是sentence-transformers配置中prompts字典的键。

例如，如果default_prompt_name为"query"，且prompts为{"query": "query: ", ...}，则句子"What is the capital of France?"将被编码为"query: What is the capital of France?"，因为提示文本将被 prepend（前置）到任何要编码的文本之前。

参数'--default-prompt-name'不能与'--default-prompt'一起使用

[env: DEFAULT_PROMPT_NAME=]

--default-prompt
默认用于编码的提示文本。如果未设置，则不应用任何提示。

例如，如果default_prompt为"query: "，则句子"What is the capital of France?"将被编码为"query: What is the capital of France?"，因为提示文本将被 prepend（前置）到任何要编码的文本之前。

参数'--default-prompt'不能与'--default-prompt-name'一起使用

[env: DEFAULT_PROMPT=]

--dense-path
可选地，定义某些嵌入模型所需的Dense模块的路径。

一些嵌入模型需要额外的Dense模块，该模块包含单个线性层和一个激活函数。默认情况下，这些Dense模块存储在2_Dense目录下，但在某些情况下，可能会提供不同的Dense模块，用于将池化嵌入转换为不同维度，如2_Dense_（例如https://huggingface.co/NovaSearch/stella_en_400M_v5）。

请注意，此参数是可选的，仅在没有modules.json文件时，或当您希望覆盖单个Dense模块路径时（仅在使用candle后端运行时）才需要设置。

[env: DENSE_PATH=]

--hf-token
您的Hugging Face Hub令牌。如果未设置--hf-token和HF_TOKEN，则会从$HF_HOME/token路径读取令牌（如果存在）。这确保了对私有或 gated 模型的访问，并允许更宽松的速率限制

[env: HF_TOKEN=]

--hostname
要监听的IP地址

[env: HOSTNAME=]
[default: 0.0.0.0]

-p, --port
要监听的端口

[env: PORT=]
[default: 3000]

--uds-path
某些text-embeddings-inference后端在与gRPC进行内部通信时将使用的unix socket名称

[env: UDS_PATH=]
[default: /tmp/text-embeddings-inference-server]

--huggingface-hub-cache
huggingface hub缓存的位置。例如，如果您想提供一个挂载磁盘，可以使用此参数覆盖位置

[env: HUGGINGFACE_HUB_CACHE=]

--payload-limit
Payload大小限制（以字节为单位）。默认值为2MB

[env: PAYLOAD_LIMIT=]
[default: 2000000]

--api-key
设置用于请求授权的API密钥。

默认情况下，服务器响应每个请求。设置API密钥后，请求必须设置Authorization头，并将API密钥作为Bearer令牌。

[env: API_KEY=]

--json-output
以JSON格式输出日志（对遥测有用）

[env: JSON_OUTPUT=]

--disable-spans
是否通过跨度包含日志跟踪

[env: DISABLE_SPANS=]

--otlp-endpoint
OpenTelemetry的gRPC端点。遥测数据通过gRPC以OTLP格式发送到此端点。例如http://localhost:4317

[env: OTLP_ENDPOINT=]

--otlp-service-name
OpenTelemetry的服务名称。例如text-embeddings-inference.server

[env: OTLP_SERVICE_NAME=]
[default: text-embeddings-inference.server]

--prometheus-port
要监听的Prometheus端口

[env: PROMETHEUS_PORT=]
[default: 9000]

--cors-allow-origin
gRPC服务器不使用此参数

[env: CORS_ALLOW_ORIGIN=]

-h, --help
打印帮助信息（使用'-h'查看摘要）

-V, --version
打印版本信息

使用私有或 gated 模型

您可以选择使用HF_TOKEN环境变量来配置text-embeddings-inference所使用的令牌。这使您能够访问受保护的资源。

例如：

1. 访问 https://huggingface.co/settings/tokens
2. 复制您的 CLI READ 令牌
3. 导出HF_TOKEN=

或使用 Docker：

model=
volume=$PWD/data # 与 Docker 容器共享卷以避免每次运行都下载权重
token=

docker run --gpus all -e HF_TOKEN=$token -p 8080:80 -v $volume:/data --pull always ghcr.io/huggingface/text-embeddings-inference:cuda-1.9 --model-id $model

离线部署

要在离线环境中部署 Text Embeddings Inference，需先下载权重，然后使用卷将其挂载到容器内。

例如：

# (可选) 创建 `models` 目录
mkdir models
cd models

# 确保已安装 git-lfs (https://git-lfs.com)
git lfs install
git clone https://huggingface.co/Qwen/Qwen3-Embedding-0.6B

# 将 models 目录设置为卷路径
volume=$PWD

# 将 models 目录通过卷挂载到容器内并设置模型 ID
docker run --gpus all -p 8080:80 -v $volume:/data --pull always ghcr.io/huggingface/text-embeddings-inference:cuda-1.9 --model-id /data/Qwen3-Embedding-0.6B

然后运行：

# 在x86上使用ONNX后端（推荐）
cargo install --path router -F ort
# 在x86上使用Intel后端
cargo install --path router -F mkl
# 在M1或M2上
cargo install --path router -F metal

现在可以通过以下命令在CPU上启动Text Embeddings Inference：

model=Qwen/Qwen3-Embedding-0.6B

text-embeddings-router --model-id $model --port 8080

[!NOTE] 在部分机器上，您可能还需要OpenSSL库和gcc。在Linux机器上，运行：

> sudo apt-get install libssl-dev gcc -y
>

CUDA

不支持CUDA计算能力 < 7.5 的GPU（V100、Titan V、GTX 1000系列等）。

确保已安装CUDA和NVIDIA驱动程序。设备上的NVIDIA驱动程序需要与CUDA 12.2或更高版本兼容。您还需要将NVIDIA二进制文件添加到路径中：

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

然后运行以下命令（可能需要一段时间，因为需要编译CUDA内核）：

# 在Turing GPU上（T4、RTX 2000系列等）
cargo install --path router -F candle-cuda-turing

# 在Ampere、Ada Lovelace、Hopper和Blackwell上
cargo install --path router -F candle-cuda

现在可以通过以下命令在GPU上启动Text Embeddings Inference：

model=Qwen/Qwen3-Embedding-0.6B

text-embeddings-router --model-id $model --port 8080

Docker

您可以使用Docker构建CPU容器，命令如下：

docker build -f Dockerfile .

要构建CUDA容器，您需要知道运行时将使用的GPU的计算能力，以便相应地构建镜像：

# 获取子模块依赖
git submodule update --init

# Turing示例（T4、RTX 2000系列等）
runtime_compute_cap=75

# Ampere示例（A100等）
runtime_compute_cap=80

# Ampere示例（A10等）
runtime_compute_cap=86

# Ada Lovelace示例（RTX 4000系列等）
runtime_compute_cap=89

# Hopper示例（H100等）
runtime_compute_cap=90

# Blackwell示例（B200、GB200等）
runtime_compute_cap=100

# Blackwell示例（GeForce RTX 50X0、RTX PRO 6000等）
runtime_compute_cap=120

# Blackwell GB10示例（DGX Spark）
runtime_compute_cap=121

docker build . -f Dockerfile-cuda --build-arg CUDA_COMPUTE_CAP=$runtime_compute_cap

ARM64 / aarch64

仅CPU（Apple Silicon、Ampere、Graviton）

对于没有NVIDIA GPU的ARM64主机，请使用CPU Dockerfile。推理仅在CPU核心上运行（不通过Docker支持Metal/MPS）。

docker build . -f Dockerfile-arm64 --platform=linux/arm64

ARM64上的CUDA（DGX Spark、Jetson）

对于带有NVIDIA GPU的ARM64主机，使用适当的计算能力和--platform linux/arm64构建Dockerfile-cuda：

# DGX Spark（GB10，sm_121）
docker build . -f Dockerfile-cuda \
--build-arg CUDA_COMPUTE_CAP=121 \
--platform linux/arm64

# 未来的ARM64 + Blackwell设备（sm_120）
docker build . -f Dockerfile-cuda \
--build-arg CUDA_COMPUTE_CAP=120 \
--platform linux/arm64

AMD Instinct GPU（ROCm）

TEI通过ROCm支持AMD Instinct GPU（MI200、MI300系列）。

model=BAAI/bge-base-en-v1.5
volume=$PWD/data

docker run \
--device /dev/kfd --device /dev/dri \
--group-add video \
--ipc=host \
-p 8080:80 \
-v $volume:/data \
--pull always \
ghcr.io/huggingface/text-embeddings-inference:rocm-latest \
--model-id $model --dtype bfloat16

有关完整的设置说明，请参见**https://huggingface.co/docs/text-embeddings-inference/amd_gpu**。

示例

• https://huggingface.co/learn/cookbook/automatic_embedding_tei_inference_endpoints
• https://github.com/plaggy/rag-containers