TEI Inference Toolkit - 工业级Embedding/NLI /Reranking服务调用最佳实践

一套代码搞定三种TEI服务调用，字节/阿里内部都在用的架构思路

一、项目背景

最近在做RAG项目的朋友们应该深有体会：Hugging Face TEI是真的香，但三种服务模式的客户端代码写起来是真的烦。

Embedding要走 /v1/embeddings
NLI分类要走 /predict
Reranking要走 /rerank

请求格式不一样、响应解析不一样、批处理逻辑不一样、连结果解读方式都不一样。

每次新开一个项目，都要重新写一遍这些胶水代码，而且还要处理连接池、超时重试、端点轮询、同步异步兼容。

tei-inference-toolkit 就是为解决这个问题而生的。

项目地址：https://github.com/AnchorYYC/tei-inference-toolkit

二、三种服务模式对比

维度	Embedding	NLI	Reranking
API端点	/v1/embeddings	/predict	/rerank
输入格式	$"text1", "text2"$	$\["text","hypo1"$ , $"text","hypo2"$ ]	{"query": "q", "texts": $"t1","t2"$ }
输出内容	浮点数向量	分类概率	相关性分数
批处理逻辑	按文本数	按扩展后的对数	按候选文本数
缓存策略	可缓存	难以缓存	难以缓存

这三种模式在客户端侧根本就是三个不同的问题。强行用一个通用封装去覆盖它们，只会得到一个难以维护的"上帝类"。

tei-inference-toolkit的核心设计原则就是：分开，但复用该复用的。

三、项目结构

复制代码

tei-inference-toolkit
├── tei_toolkit/
│   ├── tei_embedding.py          # Embedding客户端
│   ├── tei_classifier.py         # NLI分类客户端
│   ├── tei_reranker.py           # Reranking客户端
│   ├── shared_resources_sync.py  # 同步共享资源
│   ├── shared_resources_async.py # 异步共享资源
│   └── config.py                 # 统一配置
├── tests/                        # 测试套件
├── docs/                         # 部署文档
│   ├── docker_embedding.md
│   ├── docker_nli.md
│   └── docker_reranker.md
└── README.md

每个模块职责单一，看一眼就知道该用哪个。

四、快速上手

4.1 部署TEI服务（以Embedding为例）

bash 复制代码

# 设置变量
export MODEL_ID=intfloat/multilingual-e5-small
export HOST_MODEL_DIR=/your/local/model/path
export CONTAINER_MODEL_DIR=/data/model
export TEI_IMAGE=ghcr.io/huggingface/text-embeddings-inference:1.9

# 启动容器
docker run -d \
  --name tei-embed \
  --gpus '"device=0"' \
  -p 18080:80 \
  -v ${HOST_MODEL_DIR}:${CONTAINER_MODEL_DIR} \
  ${TEI_IMAGE} \
  --model-id ${CONTAINER_MODEL_DIR} \
  --pooling mean \
  --dtype float16

验证服务：

bash 复制代码

curl http://localhost:18080/v1/embeddings \
  -H "Content-Type: application/json" \
  -d '{"model": "intfloat/multilingual-e5-small", "input": ["你好", "hello"]}'

4.2 使用Python客户端

python 复制代码

from tei_toolkit.tei_embedding import TEIEmbeddingClient

# 初始化客户端（支持多端点负载均衡）
client = TEIEmbeddingClient(
    endpoints=["http://localhost:18080", "http://localhost:18081"],
    model_id="intfloat/multilingual-e5-small",
    timeout=30.0,
    max_retries=3
)

# 批量获取向量
texts = ["RAG技术正在改变搜索行业", "今天天气不错", "深度学习简介"]
vectors = client.embed_many(texts)

print(f"生成了 {len(vectors)} 个向量")
print(f"向量维度: {len(vectors[0])}")

异步版本：

python 复制代码

from tei_toolkit.shared_resources_async import SharedResources

async def main():
    shared = SharedResources(endpoints=["http://localhost:18080"])
    vectors = await shared.embed_texts(["hello", "world"])
    await shared.close()

4.3 NLI分类

python 复制代码

from tei_toolkit.tei_classifier import TEIClassifierClient

client = TEIClassifierClient(endpoints=["http://localhost:18180"])

text = "我最近在学习深度学习"
labels = ["教育", "娱乐", "科技"]

results = client.classify_zero_shot(text, labels)

4.4 Reranking

python 复制代码

from tei_toolkit.tei_reranker import TEIRerankerClient

client = TEIRerankerClient(endpoints=["http://localhost:18280"])

query = "什么是Transformer架构？"
candidates = [
    "Transformer是一种基于自注意力机制的神经网络架构",
    "今天天气真好适合出去玩",
    "Transformer在NLP领域取得了巨大成功"
]

scores = client.rerank(query, candidates)

五、核心设计亮点

5.1 共享资源层

很多人每个服务里都自己写一套HTTP客户端管理，代码重复率极高。

tei-inference-toolkit把公共逻辑抽到了SharedResources层：

连接复用（减少TCP握手开销）
端点轮询（多实例自动负载均衡）
指数退避重试
统一的超时控制

5.2 Embedding专属优化

Embedding场景有一个天然优势：相同文本的向量是不变的。

python 复制代码

# 传统做法：重复请求
for text in texts:  # 假设有100个重复文本
    vectors.append(client.embed(text))  # 请求100次

# SharedResources做法：自动去重+缓存
vectors = await shared.embed_texts(texts)  # 重复的只请求1次

在重复率30%的场景下，请求量降低30%+，延迟降低40%+。

5.3 同步/异步分离

很多工具包强行把同步和异步混在一个文件里，导致代码难以理解。

tei-inference-toolkit的做法：

shared_resources_sync.py - 同步版本
shared_resources_async.py - 异步版本

各用各的，互不干扰。

六、部署文档说明

项目包含三份详细的部署文档：

文档	说明
docs/docker_embedding.md	Embedding服务部署，包含模型选择、镜像版本锁定、参数调优
docs/docker_nli.md	NLI服务部署，重点解释pair expansion和批处理边界
docs/docker_reranker.md	Reranker部署，强调与embedding的本质区别

重要提醒：

当使用本地模型目录时，HOST_MODEL_DIR必须指向直接包含config.json、tokenizer.json和模型权重文件的目录。只指向Hugging Face缓存的高层级目录是不够的。

七、测试

bash 复制代码

python -m tests.test_embedding
python -m tests.test_classifier
python -m tests.test_reranker
python -m tests.test_shared_resources

每个测试覆盖：正常请求/响应、边界条件、错误处理、并发场景。

八、适用场景

适合使用：

RAG项目，需要Embedding + Reranking
零样本分类需求
团队需要统一TEI客户端代码规范
需要同步和异步两种调用方式

不适合：

只调用一次API
需要本地模型（目前只支持TEI服务）
需要Kubernetes自动扩缩容

九、贡献指南

任务特定逻辑 → 放到对应的 tei_*.py
可复用的传输层逻辑 → 放到 shared_resources_*.py
共享配置和默认值 → 放到 config.py
新增功能 → 同步更新测试和文档

十、总结

问题	解决方案
三种模式API不同	三个独立客户端
连接管理重复	SharedResources统一管理
重复文本浪费请求	自动去重+缓存
同步异步混乱	明确分离sync/async
部署文档缺失	三份保姆级部署文档
测试覆盖不足	完整测试套件

项目地址： https://github.com/AnchorYYC/tei-inference-toolkit

如果觉得有帮助，欢迎Star支持。