PostgreSQL和OpenAI嵌入式的语义搜索

PostgreSQL与OpenAI嵌入语义搜索的技术进展' (PostgreSQL yǔ OpenAI qiān rù yǔyì sōusuǒ de jìshù jìnzhǎn)

图像来自Unsplash上的Igor Omilaev

在企业数据库中实施语义搜索可能是具有挑战性并需要大量工作的。但是,是否一定要这样呢?在本文中,我将演示如何利用PostgreSQL和OpenAI Embeddings在数据上实现语义搜索。如果你不想使用OpenAI Embeddings API,我将为你提供免费嵌入模型的链接。

在很高的层次上,使用LLMs的向量数据库可以对可用数据(存储在数据库、文档等中)进行语义搜索。多亏了由传奇的Jeff Dean合著的“Efficient Estimation of Word Representations in Vector Space”(也称为“Word2Vec Paper”)论文,我们知道如何将单词表示为实值向量。单词嵌入是单词在向量空间中的密集向量表示,其中具有相似含义的单词靠得更近。单词嵌入捕捉单词之间的语义关系,而创建它们的技术不止一种。

图像由作者提供

让我们练习并使用OpenAI的text-embedding-ada模型!选择距离函数通常并不重要。OpenAI推荐使用余弦相似度。如果你不想使用OpenAI嵌入并且更喜欢在本地运行不同的模型而不是进行API调用,我建议考虑使用SentenceTransformers预训练模型之一。明智地选择你的模型。

import osimport openaifrom openai.embeddings_utils import cosine_similarityopenai.api_key = os.getenv("OPENAI_API_KEY")def get_embedding(text: str) -> list: response = openai.Embedding.create(     input=text,     model="text-embedding-ada-002" ) return response['data'][0]['embedding']good_ride = "good ride"good_ride_embedding = get_embedding(good_ride)print(good_ride_embedding)# [0.0010935445316135883, -0.01159335020929575, 0.014949149452149868, -0.029251709580421448, -0.022591838613152504, 0.006514389533549547, -0.014793967828154564, -0.048364896327257156, -0.006336577236652374, -0.027027441188693047, ...]len(good_ride_embedding)# 1536

现在我们已经对嵌入是什么有了了解,让我们利用它来对一些评论进行排序。

good_ride_review_1 = "我真的很喜欢这次旅行!路程非常平顺,接送地点很方便,下车点就在咖啡店前面。"good_ride_review_1_embedding = get_embedding(good_ride_review_1)cosine_similarity(good_ride_review_1_embedding, good_ride_embedding)# 0.8300454513797334good_ride_review_2 = "驾驶非常舒适。在整个旅程中,我感到很安全,还非常喜欢车上的娱乐设施,在车在行驶过程中让我有些乐趣。"good_ride_review_2_embedding = get_embedding(good_ride_review_2)cosine_similarity(good_ride_review_2_embedding, good_ride_embedding)# 0.821774476808789bad_ride_review = "在十字路口突然急刹车,我感到非常惊讶和紧张。我没有准备好。此外,我注意到车厢里有上一个乘客留下的一些垃圾。"bad_ride_review_embedding = get_embedding(bad_ride_review)cosine_similarity(bad_ride_review_embedding, good_ride_embedding)# 0.7950041130579355

尽管绝对差异可能看起来很小,但考虑到具有成千上万条评论的排序功能。在这种情况下,我们可以优先突出显示顶部的积极评论。

一旦单词或文档被转换为嵌入,它可以存储在数据库中。然而,这个操作不会自动将数据库分类为向量数据库。只有当数据库开始支持对向量的快速操作时,我们才能正当地将其标记为向量数据库。

有许多商业和开源向量数据库,这是一个非常受讨论的话题。我将使用一个开源的PostgreSQL扩展pgvector来演示向量数据库的功能,它为可能是最受欢迎的数据库启用了向量相似性搜索功能。

让我们使用pgvector运行PostgreSQL容器:

docker pull ankane/pgvectordocker run --env "POSTGRES_PASSWORD=postgres" --name "postgres-with-pgvector" --publish 5432:5432 --detach  ankane/pgvector

让我们启动pgcli以连接到数据库(pgcli postgres://postgres:postgres@localhost:5432),然后创建一个表,插入我们上面计算的嵌入,然后选择相似的项目:

-- 启用pgvector扩展CREATE EXTENSION vector;-- 创建一个具有1536维度的向量列。-- `text-embedding-ada-002`模型有1536维度.CREATE TABLE reviews (text TEXT, embedding vector(1536));-- 插入上述的三个评论。我省略了输入以方便你的使用。INSERT INTO reviews (text, embedding) VALUES ('I really enjoyed the trip! The ride was incredibly smooth, the pick-up location was convenient, and the drop-off point was right in front of the coffee shop.', '[-0.00533589581027627, -0.01026702206581831, 0.021472081542015076, -0.04132508486509323, ...');INSERT INTO reviews (text, embedding) VALUES ('The drive was exceptionally comfortable. I felt secure throughout the journey and greatly appreciated the on-board entertainment, which allowed me to have some fun while the car was in motion.', '[0.0001858668401837349, -0.004922827705740929, 0.012813017703592777, -0.041855424642562866, ...');INSERT INTO reviews (text, embedding) VALUES ('A sudden hard brake at the intersection really caught me off guard and stressed me out. I was not prepared for it. Additionally, I noticed some trash left in the cabin from a previous rider.', '[0.00191772251855582, -0.004589076619595289, 0.004269456025213003, -0.0225954819470644, ...');-- 检查一下select count(1) from reviews;-- +-------+-- | count |-- |-------|-- | 3     |-- +-------+

现在我们准备好搜索相似的文档了。我再次缩短了“好的旅程”嵌入,因为打印1536维度是过多的。

--- 这里我们使用的嵌入是“好的旅程”的SELECT substring(text, 0, 80) FROM reviews ORDER BY embedding <-> '[0.0010935445316135883, -0.01159335020929575, 0.014949149452149868, -0.029251709580421448, ...';-- +--------------------------------------------------------------------------+-- | substring                                                                |-- |--------------------------------------------------------------------------|-- | I really enjoyed the trip! The ride was incredibly smooth, the pick-u... |-- | The drive was exceptionally comfortable. I felt secure throughout the... |-- | A sudden hard brake at the intersection really caught me off guard an... |-- +--------------------------------------------------------------------------+SELECT 3Time: 0.024s

完成!正如你所观察到的,我们计算了多个文档的嵌入,将它们存储在数据库中,并进行了向量相似性搜索。潜在的应用非常广泛,从企业搜索到医疗记录系统中用于识别具有相似症状的患者。此外,这种方法不仅限于文本;还可以计算其他类型的数据,如声音、视频和图像的相似性。

尽情享受吧!