应用场景

TEI适配昇腾环境与MindIE Torch组件整体方案介绍

原生TEI仅支持GPU硬件环境，且Python后端接口仅支持Embedding模型的embed服务。为了将其适配到昇腾环境基于MindIE Torch组件进行开发，并扩展其Python后端功能，进行了以下适配。

• NPU设备的适配：
  • 设备设置：添加对NPU设备的支持，允许在NPU上执行模型推理。通过环境变量TEI_NPU_DEVICE可以指定具体的NPU设备，确保相关的编译模式和设备选项正确应用。
  • MindIE Torch支持：引入MindIE Torch组件，对Embedding和Reranker模型进行编译优化，保证模型计算能够在NPU设备上运行，并提升模型的推理性能。

• 新增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组件所支持功能特性与模型列表

• 当前已支持TEI v1.2.3版本适配昇腾环境并基于MindIE Torch组件进行开发，支持的功能如下所示：
  • embed服务：将输入文本进行分词处理，并为每个token生成对应的词嵌入(hidden_states)，然后对这些嵌入结果进行池化处理，并将池化结果作为输入文本的嵌入向量返回给请求端。对于一段分词后序列长度为N的输入文本，其返回结果为一个长度为D(嵌入模型指定隐空间特征维度)的向量。
  • embed_all服务：与embed服务类似，分词后对每个token生成对应词嵌入(hidden_states)；不同之处在于，该服务直接返回每个token的词嵌入，而不是池化后的嵌入向量。对于一段分词后序列长度为N的输入文本，其返回结果为一个维度为[N, D]的二维矩阵。
  • rerank服务：接受一个查询问题和一组文本列表作为输入，将它们分别进行分词处理，并生成每个文本与查询的相关分数，然后根据这些相关分数对文本列表进行排序，并将分数及排序结果返回给请求端。
• 支持的模型如表1所示：

  表1 模型列表

  模型类型	支持的模型
  文本嵌入模型	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

环境变量介绍