类功能

功能描述

本地使用transformers启动模型，提供文本至向量的sparse embedding功能。需要使用transformers支持的BertModel类模型权重。类继承实现了langchain_core.embeddings.Embeddings接口。当前支持的模型：bge-m3。

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

函数原型

from mx_rag.embedding.local import SparseEmbedding
SparseEmbedding(model_path, dev_id, use_fp16)

输入参数说明

参数名	数据类型	可选/必选	说明
model_path	str	必选	模型权重文件目录，路径长度不能超过1024，不能为软链接和相对路径。目录下的各文件大小不能超过10GB、深度不超过64，且文件总个数不超过512。运行用户的属组，以及非运行用户不能有该目录下文件的写权限。目录下的文件以及文件的上一级目录的属组必须是运行用户。
dev_id	int	可选	模型运行NPU id，通过npu-smi info查询可用ID，取值范围[0, 63]，默认为卡0。
use_fp16	bool	可选	是否将模型转换为半精度，默认为“True”。

参数名

数据类型

可选/必选

说明

model_path

str

必选

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

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

dev_id

int

可选

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

use_fp16

bool

可选

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

调用示例

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

父主题： SparseEmbedding