使用稀疏向量检索

import json
from typing import List

import tcvectordb
from tcvectordb.model.document import Document, AnnSearch, WeightedRerank, KeywordSearch
from tcvectordb.model.enum import FieldType, IndexType, MetricType
from tcvectordb.model.index import Index, VectorIndex, FilterIndex, HNSWParams, SparseVector
from tcvdb_text.encoder import BM25Encoder
import numpy as np

vdb_url = 'http://x.x.x.x:8100'
vdb_key = 'xB2iQyVVFy***********************'
db_name = "python-sdk2-test-sparsevector"
coll_name = "sdk2_collection-sparsevector"


# 初始化VectorDB Client
vdb_client = tcvectordb.RPCVectorDBClient(url=vdb_url,
                                          key=vdb_key,
                                          username='root')
# 初始化BM25 Encoder
bm25 = BM25Encoder.default('zh')


# 创建Database
vdb_client.drop_database(db_name)
db = vdb_client.create_database(db_name)


# 创建Collection
#  构建索引
index = Index()
index.add(FilterIndex('id', FieldType.String, IndexType.PRIMARY_KEY))
index.add(VectorIndex('vector', 768, IndexType.HNSW, MetricType.IP,
                      HNSWParams(m=16, efconstruction=200)))
index.add(FilterIndex(name='sparse_vector',
                      field_type=FieldType.SparseVector,
                      index_type=IndexType.SPARSE_INVERTED,
                      metric_type=MetricType.IP,
                      ))
#  创建Collection
coll = db.create_collection(
    name=coll_name,
    shard=1,
    replicas=0,
    description='test collection',
    index=index,
    timeout=100,
)


# 插入数据
vecs: List[List[float]] = np.random.rand(5, 768).tolist()
sparse_vectors: List[SparseVector] = bm25.encode_texts([
    '腾讯云向量数据库（Tencent Cloud VectorDB）是一款全托管的自研企业级分布式数据库服务，专用于存储、索引、检索、管理由深度神经网络或其他机器学习模型生成的大量多维嵌入向量。',
    '作为专门为处理输入向量查询而设计的数据库，它支持多种索引类型和相似度计算方法，单索引支持10亿级向量规模，高达百万级 QPS 及毫秒级查询延迟。',
    '不仅能为大模型提供外部知识库，提高大模型回答的准确性，还可广泛应用于推荐系统、NLP 服务、计算机视觉、智能客服等 AI 领域。',
    '腾讯云向量数据库（Tencent Cloud VectorDB）作为一种专门存储和检索向量数据的服务提供给用户， 在高性能、高可用、大规模、低成本、简单易用、稳定可靠等方面体现出显著优势。 ',
    '腾讯云向量数据库可以和大语言模型 LLM 配合使用。企业的私域数据在经过文本分割、向量化后，可以存储在腾讯云向量数据库中，构建起企业专属的外部知识库，从而在后续的检索任务中，为大模型提供提示信息，辅助大模型生成更加准确的答案。',
])
res = coll.upsert(documents=[
    Document(id='0001',
             vector=vecs[0],
             sparse_vector=sparse_vectors[0],
             ),
    Document(id='0002',
             vector=vecs[1],
             sparse_vector=sparse_vectors[1],
             ),
    Document(id='0003',
             vector=vecs[2],
             sparse_vector=sparse_vectors[2],
             ),
    Document(id='0004',
             vector=vecs[3],
             sparse_vector=sparse_vectors[3],
             ),
    Document(id='0005',
             vector=vecs[4],
             sparse_vector=sparse_vectors[4],
             )
])
print(res, flush=True)


# hybrid search
res = coll.hybrid_search(
    ann=[
        AnnSearch(
            field_name="vector",
            data=vecs[0],
            limit=2,
        )
    ],
    match=[
        KeywordSearch(
            field_name="sparse_vector",
            data=bm25.encode_queries('向量'),
            limit=3,
        ),
    ],
    rerank=WeightedRerank(
        field_list=['vector', 'sparse_vector'],
        weight=[0.6, 0.4],
    ),
    # rerank=RRFRerank(
    #     k=1,
    # ),
    retrieve_vector=False,
    limit=2,
)
for doc in res:
    print(json.dumps(doc, indent=2, ensure_ascii=False))


# 关闭VectorDB Client
vdb_client.drop_database(db_name)
vdb_client.close()