功能介绍

总览

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方法返回正确的输出。

• 客户端改动：

  客户端扩展：为客户端增加对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运行的模型；使用前需通过MindIE Torch对模型进行编译，当前已验证模型列表如右方所示。	文本嵌入模型	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

环境变量介绍