应用场景
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方法返回正确的输出。
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所示：