项目介绍:Text Embeddings Inference
Text Embeddings Inference(简称TEI)是一个用于部署和服务开源文本嵌入及序列分类模型的工具包。它能够高效地提取流行模型,如FlagEmbedding、Ember、GTE和E5的嵌入,实现了快速推断等多项功能。
项目特点
- 无需模型图编译:此工具包无需进行模型图编译,简化了使用步骤。
- 本地执行支持:支持在Mac设备上通过Metal进行本地执行。
- 轻量级Docker镜像:启动速度快,适合真正意义上的无服务器环境。
- 基于Token的动态批处理:采用动态批处理,优化计算效率。
- 优化推断代码:通过Flash Attention、Candle和cuBLASLt对Transformers进行推断优化。
- Safetensors权重加载:利用Safetensors进行权重加载,提升安全性和效率。
- 工业级准备:支持Open Telemetry分布式追踪与Prometheus指标,满足生产环境需要。
支持的模型
文本嵌入模型
TEI目前支持包括Nomic、BERT、CamemBERT、XLM-RoBERTa等多种模型。这些模型在不同体量和类型中表现优异,例如:
- Mistral(Salesforce/SFR-Embedding-2_R)
- Qwen2(Alibaba-NLP/gte-Qwen2-7B-instruct)
- Bert(WhereIsAI/UAE-Large-V1)
- XLM-RoBERTa(intfloat/multilingual-e5-large-instruct)
更多关于最佳文本嵌入模型的详细信息可以参阅MTEB排行榜。
序列分类与再排序模型
TEI还支持CamemBERT和XLM-RoBERTa的序列分类模型,包括用于情感分析和再排序任务的模型,例如:
- 再排序:XLM-RoBERTa(BAAI/bge-reranker-large)
- 情感分析:RoBERTa(SamLowe/roberta-base-go_emotions)
使用指南
Docker使用
用户可以通过Docker轻松部署模型。以下是一个简单的使用实例:
model=BAAI/bge-large-en-v1.5
volume=$PWD/data
docker run --gpus all -p 8080:80 -v $volume:/data --pull always ghcr.io/huggingface/text-embeddings-inference:1.5 --model-id $model
注意:要使用GPU,需要安装NVIDIA Container Toolkit,并确保NVIDIA驱动程序与CUDA 12.2版本或更高版本兼容。
本地安装
用户也可以选择在本地安装TEI:
- 安装Rust:
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
- 根据机器架构安装TEI:
- 在x86上:
cargo install --path router -F mkl
- 在M1或M2上:
cargo install --path router -F metal
- 在x86上:
API文档
详细的API文档可以通过/docs
路径访问,或者在线查看Swagger UI。
示例与更多信息
TEI提供了一系列示例来帮助用户更好地理解和应用此工具包,例如:
- 如何通过TEI设置推断端点
- 如何在RAG容器中使用TEI
Text Embeddings Inference通过其高效的部署能力和多种支持功能,为用户提供了一个强大且灵活的文本处理解决方案。无论是在科研环境还是生产应用中,TEI都能以其出色的性能和易用性,成为文本嵌入和分类任务的理想选择。