PostgreSQL - pgvector 插件构建向量数据库并举行相似度查询 ...

在现代的机器学习和人工智能应用中，向量相似度检索是一个非常紧张的技能，尤其是在文本、图像或其他类型的嵌入向量的操作中。本文将先容如何在 PostgreSQL 中安装 pgvector 插件，用于存储和检索向量数据，并展示如何通过 Python 脚本向数据库插入向量并实行相似度查询。
一、安装 PostgreSQL 并配置 pgvector 插件

1. 安装 PostgreSQL

首先，确保你已经安装了 PostgreSQL。可以通过 Docker 快速安装 PostgreSQL。

1. docker pull postgres
2. docker run --name my_postgres -e POSTGRES_PASSWORD=mysecretpassword -e POSTGRES_USER=myuser -e POSTGRES_DB=mydb -p 5432:5432 -d postgres

复制代码

上述下令会启动一个名为 my_postgres 的容器，并袒露 5432 端口以便外部连接。
2. 安装 pgvector 插件

pgvector 插件可以让你在 PostgreSQL 中存储向量，并支持高效的向量相似度查询。通过以下步骤安装 pgvector：
进入 PostgreSQL 容器后，实行以下下令：

1. docker exec -it my_postgres bash
2. apt update
3. apt install -y postgresql-server-dev-all make gcc
4. git clone https://github.com/pgvector/pgvector.git
5. cd pgvector
6. make
7. make install

复制代码

3. 创建数据库并启用 pgvector 插件

进入 PostgreSQL 终端，创建一个数据库并启用 pgvector 插件：

1. psql -U myuser -d mydb

复制代码

在 PostgreSQL 终端中实行：

1. CREATE EXTENSION vector;

复制代码

如许我们就可以使用 pgvector 插件来存储和检索向量了。
二、创建向量表

接下来，我们必要创建一个用于存储向量的表。假设我们有一个名为 knowledge.vector_data 的表，它将存储用户 ID、文件 ID、文本内容以及对应的向量。

1. -- 创建 schema
2. CREATE SCHEMA knowledge;
4. -- 创建 vector_data 表
5. CREATE TABLE knowledge.vector_data (
6.     id SERIAL PRIMARY KEY,  -- 自动递增的主键
7.     user_id BIGINT NOT NULL DEFAULT 0,
8.     file_id BIGINT NOT NULL DEFAULT 0,
9.     content VARCHAR(65535) NOT NULL DEFAULT '',  -- 存储文本内容
10.     featrue VECTOR(1024),  -- 存储1024维向量
11.     created TIMESTAMP DEFAULT CURRENT_TIMESTAMP
12. );
14. -- 创建 HNSW 索引并指定操作符类 (例如，使用余弦相似度)
15. CREATE INDEX ON knowledge.vector_data USING hnsw (featrue vector_cosine_ops);

复制代码

这个表将存储文本内容及其对应的 1024 维向量。我们选择 1024 维是由于一些嵌入模型生成的向量维度较高。如果你的向量维度差别，可以调整此值。
三、Python 测试程序：向量化并插入数据库

接下来，我们将编写一个 Python 脚本，将一些测试句子向量化后插入数据库，并举行相似度查询。
1. 安装所需 Python 库

首先，你必要安装 sentence-transformers 库来举行句子的向量化处理，并使用 psycopg2 连接 PostgreSQL：

1. pip install sentence-transformers psycopg2 numpy

复制代码

2. 完整的 Python 脚本

以下是完整的 Python 脚本，它可以将测试语句向量化后插入 PostgreSQL 数据库，然后实行相似度查询并返回相似的句子及其相似度评分：

1. import psycopg2
2. from sentence_transformers import SentenceTransformer
3. import numpy as np
5. # 数据库连接信息
6. DB_HOST = "localhost"
7. DB_PORT = "5432"
8. DB_NAME = "mydb"
9. DB_USER = "myuser"
10. DB_PASSWORD = "mysecretpassword"
12. # 连接 PostgreSQL 数据库
13. def connect_db():
14.     conn = psycopg2.connect(
15.         host=DB_HOST,
16.         port=DB_PORT,
17.         dbname=DB_NAME,
18.         user=DB_USER,
19.         password=DB_PASSWORD
20.     )
21.     return conn
23. # 向量化函数
24. def vectorize_sentences(sentences, target_dim=1024):
25.     # 使用预训练模型 "all-MiniLM-L6-v2" 向量化句子
26.     model = SentenceTransformer('paraphrase-MiniLM-L12-v2')
27.     sentence_vectors = model.encode(sentences)
29.     # 将 384 维向量扩展到 1024 维，使用 0 进行填充
30.     padded_vectors = [np.pad(vector, (0, target_dim - len(vector)), 'constant') for vector in sentence_vectors]
31.     return padded_vectors
33. # 将向量插入数据库
34. def insert_vectors_to_db(sentences, vectors):
35.     conn = connect_db()
36.     cursor = conn.cursor()
38.     # 将句子和其对应的向量插入表中
39.     for sentence, vector in zip(sentences, vectors):
40.         sql = """
41.         INSERT INTO knowledge.vector_data (user_id, file_id, content, featrue)
42.         VALUES (%s, %s, %s, %s)
43.         """
44.         cursor.execute(sql, (1, 1, sentence, vector.tolist()))  # 假设 user_id 和 file_id 为 1
45.         print(f"Inserted sentence: {sentence}")
47.     conn.commit()
48.     cursor.close()
49.     conn.close()
52. # 查询最相似的句子
53. def query_similar_sentences(target_vector):
54.     conn = connect_db()
55.     cursor = conn.cursor()
57.     # 转换向量为 PostgreSQL 可接受的格式
58.     vector_str = '[' + ','.join(map(str, target_vector)) + ']'
60.     # 查询相似句子
61.     sql = """
62.     SELECT content, 1 - (featrue <=> %s) AS similarity
63.     FROM knowledge.vector_data
64.     ORDER BY featrue <=> %s ASC
65.     LIMIT 5;
66.     """
67.     cursor.execute(sql, (vector_str, vector_str))
68.     results = cursor.fetchall()
70.     for row in results:
71.         print(f"Sentence: {row[0]}, Similarity: {row[1]}")
73.     cursor.close()
74.     conn.close()
77. if __name__ == "__main__":
78.     # 测试语句
79.     sentences = [
80.         "This is a test sentence.",
81.         "Another sentence for testing.",
82.         "PostgreSQL vector database integration."
83.     ]
85.     # 向量化测试语句
86.     # vectors = vectorize_sentences(sentences)
88.     # 将向量插入数据库
89.     # insert_vectors_to_db(sentences, vectors)
91.     # 目标句子，用于查询相似句子
92.     target_sentence = "This is a sentence."
94.     # 获取目标句子的向量
95.     target_vector = vectorize_sentences([target_sentence], target_dim=1024)[0]
97.     # 查询最相似的句子
98.     query_similar_sentences(target_vector)

复制代码

四、实行流程

• 向量化并插入数据库：通过 insert_vectors_to_db 函数，脚本会将输入的测试语句向量化，并将其存储到 PostgreSQL 中的 knowledge.vector_data 表。
  
• 相似度查询：通过 query_similar_sentences 函数，脚本会根据输入的目标句子，查询出最相似的句子，并按相似度排序返回结果。
  

五、结论

通过本文的步骤，你可以在 PostgreSQL 中使用 pgvector 插件存储和查询向量数据。这种方法可以用于文本、图像等数据的相似度检索。我们利用 sentence-transformers 生成句子嵌入，将其存储在 PostgreSQL 数据库中，并通过 SQL 查询高效地返回相似的句子。
使用向量数据库，可以让我们在处理大量嵌入数据时，充分发挥数据库和机器学习的联合力量，实现高效、快速的向量检索。

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。