昇腾社区首页
中文
注册

类功能

功能描述

用于配置similarity cache的配置数据结构。

函数原型

from mx_rag.cache import SimilarityCacheConfig
SimilarityCacheConfig(self,vector_config, cache_config,emb_config,similarity_config,retrieval_top_k,clean_size,**kwargs)

输入参数说明

参数名

数据类型

可选/必选

说明

vector_config

Dict[str, Any]

必选

配置矢量数据库。具体配置参见表1,注意这里如果是faiss数据库,MindFAISS的load_local_index参数会被data_save_folder落盘路径覆盖,auto_save参数置为False,配置中的字符串长度不能超过1024,字典中包含的可迭代序列长度不能超过1024,字典长度不能超过1024,字典嵌套深度不能超过2层。

cache_config

str

必选

配置标量数据库。当前只支持配置为“sqlite”。

emb_config

Dict[str, Any]

必选

配置embedding模型,请参见表2。字典长度不能超过1024,字典中字符串长度不能超过1024,字典嵌套深度不能超过1层。

similarity_config

Dict[str, Any]

必选

配置相似度计算模型,字典长度不能超过1024,字典中字符串长度不能超过1024,字典嵌套深度不能超过1层。

请参考表3

retrieval_top_k

int

可选

相似检索时的topk值。默认值为“1”

取值范围:(0, 1000]。

clean_size

int

可选

每次缓存数据添加超过cache_size时,老化的个数,默认值为“1”

取值范围:(0, cache_size]

**kwargs

Any

必选

参数介绍可参见CacheConfig

  • 本接口内部使用了pickle模块,有被恶意构造的数据在unpickle期间攻击的风险。需要保证在被加载的落盘数据data_save_folder是安全存储,仅可加载可信的落盘数据。
  • vector_config和cache_config必须同时为none或同时不为none。如果vector_config和cache_config同时为None,则等同于memory cache。
  • 对于sqlite数据库来说,它的落盘文件不能超过30GB,对于矢量数据库来说,它的落盘文件不能超过20GB。
表1 vector_config

参数名

数据类型

可选/必选

说明

**kwargs

Dict[str, Any]

必选

具体介绍可参见create_storage

top_k

int

可选

相似度检索时的topk个数。默认值:5。

vector_save_file

str

必选

落盘路径。vector_type为"npu_faiss_db"时,该参数会覆盖MindFAISS中的load_local_index参数作为落盘路径;对于“milvus_db”,该参数不生效。

表2 emb_config参数说明

参数名

数据类型

可选/必选

说明

x_dim

int

可选

embedding模型的维度,默认值为“0”

skip_emb

bool

可选

是否跳过embedding,默认值为“False”

**kwargs

Dict[str, Any]

必选

具体介绍可参见create_embedding

表3 similarity_config参数说明

参数名

数据类型

可选/必选

说明

score_min

float

可选

相似度计算可能存在范围的最小值,默认值为“0.0”

取值范围:[0.0, 100.0]

score_max

float

可选

相似度计算可能存在范围的最大值,默认值为“1”

取值范围:[1.0, 100.0],“score_max”必须大于等于“score_min”

reverse

bool

可选

相似度分数和相似度的关系,默认值为“False”

  • False:相似度分数越高,相似度越高。
  • True:相似度分数越高,相似度越低。

**kwargs

Dict[str, Any]

必选

具体介绍可参见create_reranker

调用示例

示例1:faiss_npu + local_embedding + local_reranker
from mx_rag.cache import SimilarityCacheConfig
from mx_rag.cache import MxRAGCache
dim = 1024
dev = 1

similarity_config = SimilarityCacheConfig(
        vector_config={
            "vector_type": "npu_faiss_db",
            "x_dim": dim,
            "devs": [dev],
           
        },
        cache_config="sqlite",
        emb_config={
            "embedding_type": "local_text_embedding",
            "x_dim": dim,
            "model_path": "path_to_embedding_model", # emb 模型路径
            "dev_id": dev
        },
        similarity_config={
            "similarity_type": "local_reranker",
            "model_path": "path_to_reranker_model",  # reranker 模型路径
            "dev_id": dev
        },
        retrieval_top_k=1,
        cache_size=1000,
        clean_size=20,
        similarity_threshold=0.86,
        data_save_folder="path_to_cache_save_folder", # 落盘路径
        disable_report=True
    )
similarity_cache = MxRAGCache("similarity_cache", similarity_config)

示例2:milvus_db + tei_embedding + tei_reranker

from paddle.base import libpaddle
from mx_rag.cache import SimilarityCacheConfig
from mx_rag.cache import EvictPolicy
from mx_rag.cache import MxRAGCache
from mx_rag.utils import ClientParam
from pymilvus import MilvusClient
dim = 1024

client = MilvusClient("https://x.x.x.x:port")
similarity_config = SimilarityCacheConfig(
    vector_config={
        "client": client,
        "vector_type": "milvus_db",
        "x_dim": dim,
       
        "collection_name": "mxrag_cache_123",  # milvus db的标签
        "param": None
    },
    cache_config="sqlite",
    emb_config={
        "embedding_type": "tei_embedding",
        "url": "https://<ip>:<port>/embed",  # tei_embedding 服务的IP地址和侦听端口
        "client_param": ClientParam(ca_file="/path/to/ca.crt")
    },
    similarity_config={
        "similarity_type": "tei_reranker",
        "url": "https://<ip>:<port>/rerank",  # tei_reranker 服务的IP地址和侦听端口
        "client_param": ClientParam(ca_file="/path/to/ca.crt")
    },
    retrieval_top_k=1,
    cache_size=100,
    auto_flush=100,
    similarity_threshold=0.70,
    data_save_folder="path_to_cache_save_folder",
    disable_report=True,
    eviction_policy=EvictPolicy.FIFO
)
similarity_cache = MxRAGCache("similarity_cache", similarity_config)