随着人工智能和深度学习技术的迅猛发展，矢量数据在数据科学和应用领域中扮演着越来越重要的角色。从自然语言处理到图像识别，矢量数据是我们理解和处理各种信息的关键。然而，有效地存储和查询矢量数据一直是一个挑战，特别是在云计算环境中。在本教程中，我们将介绍如何使用Neon Serverless Postgres构建OpenAI矢量数据库，以便存储和查询OpenAI生成的嵌入向量。
准备工作
在开始之前，请确保您具备以下条件：
- 一个Neon Postgres数据库。您可以在几个简单的步骤中创建一个带有现成的neondb数据库的帐户并设置项目。有关说明，请参阅注册并创建您的第一个项目。
- 您的Neon数据库连接字符串。您可以从Neon仪表板上的连接详细信息小部件中复制它。请参阅从任何应用程序连接。
- pgvector扩展。在Neon中运行CREATE EXTENSION vector;来安装扩展。有关说明，请参阅启用pgvector扩展。
- 您的OpenAI API密钥。
- Python和pip。
安装所需模块
这个笔记本需要openai、psycopg2、pandas、wget和python-dotenv模块。您可以使用pip来安装它们：
!pip install openai psycopg2 pandas wget python-dotenv准备您的OpenAI API密钥
为了生成文档和查询的向量，需要一个OpenAI API密钥。
如果您没有OpenAI API密钥，请从https://platform.openai.com/account/api-keys获取一个。
将OpenAI API密钥添加为操作系统环境变量或在会话中提供它。如果您定义了一个环境变量，请将变量命名为OPENAI_API_KEY。
测试您的OpenAI密钥
import os
from getpass import getpass
# 检查是否将OPENAI_API_KEY设置为环境变量
if os.getenv("OPENAI_API_KEY") is not None:
print("您的OPENAI_API_KEY已准备就绪")
else:
# 如果没有，请提示输入
api_key = getpass("请输入您的OPENAI_API_KEY: ")
if api_key:
print("您的OPENAI_API_KEY现在可用于此会话")
# 可选地，您可以将其设置为当前会话的环境变量
os.environ["OPENAI_API_KEY"] = api_key
else:
print("您未输入OPENAI_API_KEY")连接到您的Neon数据库
请在下面提供您的Neon数据库连接字符串，或者使用DATABASE_URL变量在.env文件中定义它。有关获取Neon连接字符串的信息，请参阅从任何应用程序连接。
import os
import psycopg2
from dotenv import load_dotenv
# 从.env文件加载环境变量
load_dotenv()
# 连接字符串可以直接在这里提供。
# 用您的Neon连接字符串替换下一行。
connection_string = "postgres://<user>:<password>@<hostname>/<dbname>"
# 如果上面没有直接提供connection_string，
# 则检查环境或.env中是否设置了DATABASE_URL。
if not connection_string:
connection_string = os.environ.get("DATABASE_URL")
# 如果两种方法都没有提供连接字符串，则引发错误。
if not connection_string:
raise ValueError("请在代码中或.env文件中提供有效的连接字符串作为DATABASE_URL。")
# 使用连接字符串连接
connection = psycopg2.connect(connection_string)
# 创建新的游标对象
cursor = connection.cursor()
# 测试数据库连接
cursor.execute("SELECT 1;")
result = cursor.fetchone()
# 检查查询结果
if result == (1,):
print("您的数据库连接成功！")
else:
print("连接失败。")加载数据
将预先计算的矢量数据加载到名为articles的数据库表中，从.csv文件加载数据。由于有25000条记录，因此预计操作将花费几分钟。
import io
# 本地CSV文件的路径
csv_file_path = '../../data/vector_database_wikipedia_articles_embedded.csv'
# 定义一个生成器函数来处理csv文件
def process_file(file_path):
with open(file_path, 'r', encoding='utf-8') as file:
for line in file:
yield line
# 创建一个StringIO对象来存储修改后的行
modified_lines = io.StringIO(''.join(list(process_file(csv_file_path))))
# 创建用于copy_expert的COPY命令
copy_command = '''
COPY public.articles (id, url, title, content, title_vector, content_vector, vector_id)
FROM STDIN WITH (FORMAT CSV, HEADER true, DELIMITER ',');
'''
# 使用copy_expert执行COPY命令
cursor.copy_expert(copy_command, modified_lines)
# 提交更改
connection.commit()
# 检查记录数以确保数据已加载。应该有25000条记录。
count_sql = """select count(*) from public.articles;"""
cursor.execute(count_sql)
result = cursor.fetchone()
print(f"记录数:{result[0]}")搜索您的数据
在将数据存储在Neon数据库中后，您可以查询最近的邻居。
首先，定义query_neon函数，当运行矢量相似性搜索时，该函数将执行。该函数基于用户的查询创建一个嵌入向量，准备SQL查询，并使用嵌入向量运行SQL查询。您加载到数据库中的预计算嵌入是使用text-embedding-ada-002 OpenAI模型创建的，因此您必须使用相同的模型来创建用于相似性搜索的嵌入。
提供了一个vector_name参数，允许基于“标题”或“内容”搜索。
def query_neon(query, collection
_name, vector_name="title_vector", top_k=20):
# 从用户查询创建嵌入向量
embedded_query = openai.Embedding.create(
input=query,
model="text-embedding-ada-002",
)["data"][0]["embedding"]
# 将嵌入查询转换为PostgreSQL兼容格式
embedded_query_pg = "[" + ",".join(map(str, embedded_query)) + "]"
# 创建SQL查询
query_sql = f"""
SELECT id, url, title, l2_distance({vector_name},'{embedded_query_pg}'::VECTOR(1536)) AS similarity
FROM {collection_name}
ORDER BY {vector_name} <-> '{embedded_query_pg}'::VECTOR(1536)
LIMIT {top_k};
"""
# 执行查询
cursor.execute(query_sql)
results = cursor.fetchall()
return results
# 基于title_vector嵌入运行相似性搜索
query_results = query_neon("Greek mythology", "Articles")
for i, result in enumerate(query_results):
print(f"{i + 1}. {result[2]} (分数: {round(1 - result[3], 3)})")
# 基于content_vector嵌入运行相似性搜索
query_results = query_neon("Famous battles in Greek history", "Articles", "content_vector")
for i, result in enumerate(query_results):
print(f"{i + 1}. {result[2]} (分数: {round(1 - result[3], 3)})")结论
通过使用Neon Serverless Postgres构建OpenAI矢量数据库，您可以有效地存储和查询OpenAI生成的嵌入向量。这为您提供了强大的工具来处理和分析文本数据，从而使您能够更好地理解和利用矢量数据的潜力。
试用ChatGpt-plus新功能：DALL·E 3
试用ChatGpt-plus新功能：DALL·E 3
试用ChatGpt-plus新功能：DALL·E 3