功能介绍
总览
TEI（全称：Text Embeddings Inference）是由Huggingface推出的高性能推理框架，旨在简化和加速文本嵌入（Embedding）和重排序（Reranker）模型在生产环境中的部署。
TEI支持基于HTTP和gRPC协议的服务接口，能够高效处理文本嵌入的生成和基于文本相关性的重排序等任务；TEI框架同时也支持多种嵌入模型和重排序模型，并提供了灵活的请求批处理、模型管理和动态调度功能。通过TEI，开发者可以轻松地扩展和优化文本嵌入和重排序服务，以满足实时应用和批量处理的需求，特别适用于自然语言处理相关的在线推理任务，能够满足RAG（全称：Retrieval-Augmented Generation）、信息检索 (IR)、自然语言理解 (NLU)、文本分类以及个性化推荐系统等下游应用场景。以下内容介绍了如何面向昇腾推理环境适配TEI框架，并拉起经过MindIE Torch与ATB后端组图优化后的模型。
TEI适配昇腾环境及接入多种模型后端（MindIE Torch与ATB）整体方案介绍
原生TEI仅支持GPU硬件环境，且Python后端接口仅支持Embedding模型的embed服务。为拓展其Python后端功能、将其适配到昇腾环境，并能够基于MindIE Torch与ATB进行组图优化，进行了以下适配。
- NPU设备的适配：
- 设备设置：添加对NPU设备的支持，允许在NPU上执行模型推理。通过环境变量TEI_NPU_DEVICE与TEI_NPU_BACKEND可以指定具体的NPU设备号与模型组图后端，确保相关的编译模式和设备选项正确应用。
- MindIE Torch组图支持：引入MindIE Torch组件，对Embedding和Reranker模型进行编译优化，保证模型计算能够在NPU设备上运行，并提升模型的推理性能。
- ATB组图支持：引入ATB加速库，通过调用昇腾Transformer加速库算子实现Embedding和Reranker模型的整网组图，提升模型的推理性能。
- 新增embed_all接口：
- 支持返回所有token的embedding：增加embed_all接口，用于返回句子中每个token的完整hidden_state，而不仅仅是返回句子中首个token的文本嵌入。
- gRPC 扩展：在gRPC层面添加了对embed_all接口的支持，修改proto文件以支持新的返回类型TokenEmbedding，并确保客户端可以调用并接收返回结果。
- 新增rerank服务：
- 引入Reranker模型：增加Reranker模型的支持，允许使用特定模型对查询和候选文本进行重新排序。新增了predict接口以处理重排序请求，并返回预测分数。
- gRPC扩展：为rerank服务扩展gRPC接口，修改服务端逻辑以支持predict接口的调用，确保重排序模型能够通过gRPC提供服务。
- 扩展和优化Python后端：
- 模型选择逻辑调整：在模型加载逻辑中为Embedding和Reranker模型增加了适配条件，确保正确加载模型到NPU设备或其他设备。
- DefaultModel与RerankModel类适配：修改嵌入模型和重排序模型的推理过程，确保在NPU设备上能够正常执行，并通过embed、embed_all和predict方法返回正确的输出。
TEI框架接入MindIE Torch与ATB后端所支持功能特性与模型列表
- 当前已支持TEI v1.2.3版本适配昇腾环境并通过MindIE Torch与ATB组件进行模型整网组图，支持的功能如下所示：
- embed服务：将输入文本进行分词处理，并为每个token生成对应的词嵌入(hidden_states)，然后对这些嵌入结果进行池化处理，并将池化结果作为输入文本的嵌入向量返回给请求端。对于一段分词后序列长度为N的输入文本，其返回结果为一个长度为D(嵌入模型指定隐空间特征维度)的向量。
- embed_all服务：与embed服务类似，分词后对每个token生成对应词嵌入(hidden_states)；不同之处在于，该服务直接返回每个token的词嵌入，而不是池化后的嵌入向量。对于一段分词后序列长度为N的输入文本，其返回结果为一个维度为[N, D]的二维矩阵。
- rerank服务：接受一个查询问题和一组文本列表作为输入，将它们分别进行分词处理，并生成每个文本与查询的相关分数，然后根据这些相关分数对文本列表进行排序，并将分数及排序结果返回给请求端。
- 支持的模型如表1所示：
表1 模型列表 模型后端
支持模型范围
模型类型
支持模型参考
MindIE Torch
MindIE Torch后端当前支持大部分能基于pytorch+torch_npu运行的模型；使用前需通过mindietorch对模型进行编译，当前已验证模型列表如右方所示。
文本嵌入模型
bge-large-zh-v1.5
jina-embeddings-v2-base-en
jina-embeddings-v2-base-code
nomic-embed-text-v1.5
nomic-embed-text-v1.0
gte-large-en-v1.5
multilingual-e5-large-instruct
重排序模型
bge-reranker-large
ATB
ATB后端当前支持所有基于Bert或XLM-Roberta结构的文本嵌入模型（例如：bge系列、bce系列）以及gte-Qwen2-1.5B-instruct；替换权重路径中的模型类文件（例如：modeling_qwen.py）后方可直接运行，无需编译。
文本嵌入模型
bge-large-zh-v1.5
bge-m3
gte-Qwen2-1.5B-instruct