使用TileDB进行多维数组的ANN搜索-优快云博客

本文链接：https://blog.youkuaiyun.com/qahaj/article/details/146465424

在现代数据处理和机器学习应用中，处理多维数组如图像、时间序列和科学数据，是一个普遍的需求。TileDB 是一个强大的引擎，专注于对稠密和稀疏多维数组的索引和查询。它不仅支持传统的数组操作，还提供了向量相似性搜索（ANN）的能力，通过 TileDB-Vector-Search 模块，你可以轻松地在本地磁盘和云对象存储（如AWS S3）上执行向量索引和无服务器ANN查询。

核心原理解析

TileDB 作为向量数据库，能够高效地执行 ANN 搜索，这对于需要快速检索相似向量的场景非常有用。通过存储向量索引，TileDB 提供了快速访问和查询能力。其向量搜索支持本地和云存储，使得在大规模数据上运行也保持高效。

代码实现演示

接下来我们将演示如何使用 TileDB 进行向量搜索。

安装必要的库

首先，确保安装了必要的 Python 包：

%pip install --upgrade --quiet tiledb-vector-search langchain-community

基本示例

在下面的示例中，我们将加载文本数据，生成嵌入，并使用 TileDB 存储和检索向量。

from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import TileDB
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 加载文本数据
raw_documents = TextLoader("../../how_to/state_of_the_union.txt").load()

# 分割文本为小块
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
documents = text_splitter.split_documents(raw_documents)

# 生成嵌入
embeddings = HuggingFaceEmbeddings()

# 创建 TileDB 实例，存储嵌入向量
db = TileDB.from_documents(
    documents, embeddings, index_uri="/tmp/tiledb_index", index_type="FLAT"
)

# 查询示例
query = "What did the president say about Ketanji Brown Jackson"
docs = db.similarity_search(query)
print(docs[0].page_content)

向量相似性搜索

你可以直接使用嵌入向量进行搜索：

# 生成查询的嵌入向量
embedding_vector = embeddings.embed_query(query)

# 使用嵌入向量进行相似性搜索
docs = db.similarity_search_by_vector(embedding_vector)
print(docs[0].page_content)

带分数的相似性搜索

如果需要检索结果的相关性分数，可以这样做：

# 获取搜索结果及其相关性分数
docs_and_scores = db.similarity_search_with_score(query)
print(docs_and_scores[0])

最大边缘相关搜索 (MMR)

除了相似性搜索，还可以使用 MMR 来提高结果的多样性：

# 使用 MMR 搜索
retriever = db.as_retriever(search_type="mmr")
retriever.invoke(query)

# 直接调用 MMR 搜索方法
db.max_marginal_relevance_search(query, k=2, fetch_k=10)