使用Scikit-learn实现向量数据库的持久化和查询

最新推荐文章于 2025-04-26 20:08:02 发布

fgayif

最新推荐文章于 2025-04-26 20:08:02 发布

阅读量352

点赞数 3

文章标签： scikit-learn 数据库 python

本文链接：https://blog.youkuaiyun.com/fgayif/article/details/146464998

版权

在机器学习项目中，我们常常需要处理大量文档，并通过相似度搜索来查找相关信息。Scikit-learn提供了强大的机器学习算法，而SKLearnVectorStore扩展了这些实现，使得我们可以将向量存储持久化为多种格式，如JSON、BSON或Apache Parquet。这篇文章将向大家展示如何使用SKLearnVectorStore进行向量数据库的持久化和查询。

技术背景介绍

Scikit-learn是一个开源的机器学习库，其中包含了一些流行的机器学习算法的实现。最近，SKLearnVectorStore将Scikit-learn的功能扩展为能够持久化向量数据库，并支持多种存储格式。这对于需要长期保存和快速访问向量数据的应用场景至关重要。

核心原理解析

SKLearnVectorStore通过使用Scikit-learn的算法实现数据的向量化和存储。它结合了OpenAI提供的嵌入功能，使得我们能够在持久化的向量数据库中进行高效的相似度搜索。使用这一工具，需要通过langchain-community进行安装，该工具是对开源机器学习功能的扩展。

代码实现演示

下面我们通过代码展示如何创建一个向量数据库并进行查询：

# 安装所需的库
%pip install --upgrade --quiet scikit-learn langchain-community bson pandas pyarrow

import os
import tempfile
from getpass import getpass
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import SKLearnVectorStore
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = getpass("Enter your OpenAI key:")

# 加载文档
loader = TextLoader("path/to/your/document.txt")
documents = loader.load()

# 文档拆分及嵌入
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
embeddings = OpenAIEmbeddings()

# 设置持久化路径
persist_path = os.path.join(tempfile.gettempdir(), "document_vectors.parquet")

# 创建向量数据库并进行查询
vector_store = SKLearnVectorStore.from_documents(
    documents=docs,
    embedding=embeddings,
    persist_path=persist_path,
    serializer="parquet"
)

query = "Specific inquiry about your document"
docs = vector_store.similarity_search(query)
print(docs[0].page_content)

# 保存向量数据库
vector_store.persist()
print("Vector store was persisted to", persist_path)

# 加载持久化的向量数据库
vector_store2 = SKLearnVectorStore(
    embedding=embeddings, persist_path=persist_path, serializer="parquet"
)
print("Loaded vector store from", persist_path)
docs = vector_store2.similarity_search(query)
print(docs[0].page_content)