Skip to main content

scikit-learn

scikit-learn 是一个开源的机器学习算法集合，包括一些 k 最近邻 的实现。SKLearnVectorStore 封装了这个实现，并添加了将向量存储为 json、bson（二进制 json）或 Apache Parquet 格式的功能。

本笔记本展示了如何使用 SKLearnVectorStore 向量数据库。

# # 如果您计划使用 bson 序列化，请还要安装：
# %pip install bson

# # 如果您计划使用 parquet 序列化，请还要安装：
%pip install pandas pyarrow

要使用 OpenAI 嵌入，您需要一个 OpenAI 密钥。您可以在 https://platform.openai.com/account/api-keys 获取一个，或者随意使用其他嵌入。

import os
from getpass import getpass

os.environ["OPENAI_API_KEY"] = getpass("输入您的 OpenAI 密钥：")

基本用法​

加载示例文档语料库​

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import SKLearnVectorStore
from langchain.document_loaders import TextLoader

loader = TextLoader("../../../state_of_the_union.txt")
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
embeddings = OpenAIEmbeddings()

创建 SKLearnVectorStore，索引文档语料库并运行示例查询​

import tempfile

persist_path = os.path.join(tempfile.gettempdir(), "union.parquet")

vector_store = SKLearnVectorStore.from_documents(
    documents=docs,
    embedding=embeddings,
    persist_path=persist_path,  # persist_path 和 serializer 是可选的
    serializer="parquet",
)

query = "总统对 Ketanji Brown Jackson 说了什么"
docs = vector_store.similarity_search(query)
print(docs[0].page_content)
    今晚。我呼吁参议院：通过《自由投票法案》。通过《约翰·刘易斯选举权法案》。顺便说一下，通过《披露法案》，这样美国人就可以知道谁在资助我们的选举了。
    
    今晚，我想向一个致力于为这个国家服务的人致敬：司法部长斯蒂芬·布雷耶——一位陆军退伍军人、宪法学者和即将退休的美国最高法院法官。布雷耶法官，感谢您的服务。
    
    总统最重要的宪法责任之一是提名人选担任美国最高法院法官。
    
    我在4天前做到了，当时我提名了巡回上诉法院法官 Ketanji Brown Jackson。她是我们国家顶级的法律智慧之一，将继续布雷耶法官的卓越传统。

保存和加载向量存储​

vector_store.persist()
print("向量存储已保存到", persist_path)
    向量存储已保存到 /var/folders/6r/wc15p6m13nl_nl_n_xfqpc5c0000gp/T/union.parquet
vector_store2 = SKLearnVectorStore(
    embedding=embeddings, persist_path=persist_path, serializer="parquet"
)
print("从", persist_path, "加载了一个新的向量存储实例")
    从 /var/folders/6r/wc15p6m13nl_nl_n_xfqpc5c0000gp/T/union.parquet 加载了一个新的向量存储实例
docs = vector_store2.similarity_search(query)
print(docs[0].page_content)
    今晚。我呼吁参议院：通过《自由投票法案》。通过《约翰·刘易斯选举权法案》。顺便说一下，通过《披露法案》，这样美国人就可以知道谁在资助我们的选举了。
    
    今晚，我想向一个致力于为这个国家服务的人致敬：司法部长斯蒂芬·布雷耶——一位陆军退伍军人、宪法学者和即将退休的美国最高法院法官。布雷耶法官，感谢您的服务。
    
    总统最重要的宪法责任之一是提名人选担任美国最高法院法官。
    
    我在4天前做到了，当时我提名了巡回上诉法院法官 Ketanji Brown Jackson。她是我们国家顶级的法律智慧之一，将继续布雷耶法官的卓越传统。

清理​

os.remove(persist_path)