昇腾社区首页
中文
注册

类功能

功能描述

本地使用transformers启动模型,提供文本至向量的embedding功能。需要使用transformers支持的BertModel类模型权重。类继承实现了langchain_core.embeddings.Embeddings接口。当前支持的模型:bge-large-zh-v1.5,acge_text_embedding。

TextEmbedding模型当前支持bert、roberta、xlm_roberta类embedding模型向量化推理加速,支持的数据类型为float16。该功能如需使用向量化推理加速,安装RAG SDK软件包时需将算子模块一并安装,并且确保已经启用该功能(默认为关闭状态),具体示例可以参考开启推理加速调用示例

配置的模型如果不是safetensors权重格式,请先将模型权重转换为safetensors格式后再使用,防止使用ckpt、bin等不安全的模型权重格式引入安全问题。

各类模型推理加速配置如下:

from modeling_bert_adapter import enable_bert_speed

from modeling_roberta_adapter import enable_roberta_speed

from modeling_xlm_roberta_adapter import enable_xlm_roberta_speed

通过设置ENABLE_BOOST变量激活模型推理加速,设置值为"True"或"False"

模型加速日志相关环境变量说明

ATB_LOG_TO_STDOUT:设置为1时表示日志记录到标准输出

ATB_LOG_TO_FILE:设置为1时表示日志记录到文件

ATB_LOG_LEVEL:设置日志等级,可配置为TRACE,DEBUG, INFO,WARN, ERROR, FATAL

函数原型

from mx_rag.embedding.local import TextEmbedding
TextEmbedding(model_path, dev_id, use_fp16, pooling_method, lock)

输入参数说明

参数名

数据类型

可选/必选

说明

model_path

str

必选

模型权重文件目录,路径长度不能超过1024,不能为软链接和相对路径。

  • 目录下的各文件大小不能超过10GB、深度不超过64,且文件总个数不超过512。
  • 运行用户的属组,以及非运行用户不能有该目录下文件的写权限。
  • 目录下的文件以及文件的上一级目录的属组必须是运行用户。

dev_id

int

可选

模型运行NPU id,通过npu-smi info查询可用ID,取值范围[0, 63],默认为卡0。

use_fp16

bool

可选

是否将模型转换为半精度,默认为“True”

pooling_method

str

可选

选择处理last_hidden_state的方式,取值范围['cls', 'mean', 'max', 'lasttoken'],默认'cls'。

lock

multiprocessing.synchronize.Lock或_thread.LockType

可选

local model不支持多线程或者多进程进行处理,如果用户需要多进程或者多线程调用此接口需要申请锁。默认值为None。

可选值:

  • None:表示不使用锁,此时该接口不支持并发。
  • multiprocessing.Lock():表示进程锁,此时该接口支持多进程调用。
  • threading.Lock():表示线程锁。此时该接口支持多线程调用。

不启用推理加速调用示例

from paddle.base import libpaddle
from mx_rag.embedding.local import TextEmbedding
# 同embed = TextEmbedding("/path/to/model", 1)
embed = TextEmbedding.create(model_path="/path/to/model", dev_id=1)
print(embed.embed_documents(['abc', 'bcd']))
print(embed.embed_query('abc'))

开启推理加速调用示例

import os
from paddle.base import libpaddle
# 适配向量化推理加速
from modeling_bert_adapter import enable_bert_speed
# 使能向量化推理加速(设置为"True"时表示使能,"False"表示不使能)
os.environ["ENABLE_BOOST"] = "True"
from mx_rag.embedding.local import TextEmbedding
# 同embed = TextEmbedding("/path/to/model", 1)
embed = TextEmbedding.create(model_path="/path/to/model", dev_id=1)
print(embed.embed_documents(['abc', 'bcd']))
print(embed.embed_query('abc'))

多线程调用示例(其余嵌入模型均可参考该示例)

from paddle.base import libpaddle
import threading
from mx_rag.embedding.local import TextEmbedding
def infer(k, embed):
    print(f"thread_{k}")
    print(embed.embed_query('abc'))
    print(embed.embed_documents(['abc', 'bcd']))

if __name__ == '__main__':
    worker_nums=2
    threads = []
    embed = TextEmbedding.create(model_path='/path/to/model', dev_id=1, pooling_method='cls', lock=threading.Lock())
    for i in range(worker_nums):
        thread = threading.Thread(target=infer, args=(i,embed,))
        threads.append(thread)
        thread.start()
    for thread in threads:
        thread.join()