用 pgvector 在 PostgreSQL 中存储 Embeddings 完全指南

Supabase 现已推出一个新的 PostgreSQL 扩展：pgvector，一个开源的向量相似性搜索工具。

过去一年 AI 功能的指数级进步激发了许多新的现实应用。一个特别的挑战是如何大规模地存储和查询 embeddings。在这篇文章中，我们将解释什么是 embeddings、为什么我们可能想使用它们，以及如何使用 pgvector 在 PostgreSQL 中存储和查询它们。

🆕 Supabase 现已发布了一个开源工具包，用于使用 Postgres 和 pgvector 开发 AI 应用。在 AI & Vectors 文档中了解更多信息。

什么是 embeddings？

Embeddings 捕捉文本、图像、视频或其他类型信息的"相关性"。这种相关性最常用于：

搜索：搜索词与文本体有多相似？

推荐：两个产品有多相似？

分类：我们如何对文本体进行分类？

聚类：我们如何识别趋势？

让我们探索一个文本 embeddings 的例子。假设我们有三个短语：

"The cat chases the mouse"

"The kitten hunts rodents"

"I like ham sandwiches"

你的任务是对意思相似的短语进行分组。如果你是人类，这应该很显而易见。短语 1 和 2 几乎相同，而短语 3 的含义完全不同。

虽然短语 1 和 2 相似，但它们没有共同的词汇（除了"the"）。然而它们的含义几乎相同。我们如何能教计算机这些是相同的？

人类使用词汇和符号来交流语言。但孤立的词汇大多是无意义的——我们需要从共享的知识和经验中获取才能理解它们。"You should Google it"这个短语只有在你知道 Google 是一个搜索引擎且人们一直在把它用作动词时才有意义。

同样，我们需要训练一个神经网络模型来理解人类语言。一个有效的模型应该在数百万个不同的例子上进行训练，以理解每个词、短语、句子或段落在不同上下文中可能的含义。

那么这与 embeddings 有什么关系呢？

Embeddings 如何工作？

Embeddings 将离散信息（词汇和符号）压缩为分布式连续值数据（向量）。如果我们把之前的短语绘制在图表上，它可能看起来像这样：

短语 1 和 2 会被绘制在彼此接近的位置，因为它们的意思相似。我们期望短语 3 位于远处，因为它不相关。如果我们有第四个短语"Sally ate Swiss cheese"，它可能存在于短语 3（奶酪可以放在三明治上）和短语 1（老鼠喜欢瑞士奶酪）之间的某处。

在这个例子中，我们只有 2 个维度：X 轴和 Y 轴。实际上，我们需要更多的维度来有效地捕捉人类语言的复杂性。

OpenAI 提供了一个 API 来使用其语言模型为一段文本生成 embeddings。你给它输入任何文本信息（博客文章、文档、你的公司知识库），它会输出一个浮点数向量，代表该文本的"含义"。

与我们上面的 2 维例子相比，他们最新的 embedding 模型 text-embedding-ada-002 将输出 1536 维。

为什么这有用？一旦我们为多个文本生成了 embeddings，使用向量数学操作（如余弦距离）计算它们的相似程度就变得轻而易举。一个完美的用例是搜索。你的流程可能看起来像这样：

预处理你的知识库并为每一页生成 embeddings

存储你的 embeddings 以供日后参考（后面会详细说明）

构建一个搜索页面，提示用户输入

获取用户输入，生成一次性 embedding，然后对你的预处理 embeddings 执行相似性搜索

向用户返回最相似的页面

Embeddings 实践

在小规模应用中，你可以将 embeddings 存储在 CSV 文件中，将其加载到 Python 中，然后使用 numPy 等库使用余弦距离或点积等方法计算它们之间的相似性。OpenAI 有一个这样做的代码示例。不幸的是，这可能不会很好地扩展：

如果我需要存储和搜索大量文档和 embeddings（超过能放入内存的数量）怎么办？

如果我想动态地创建/更新/删除 embeddings 怎么办？

如果我没有使用 Python 怎么办？

pgvector 登场，这是 PostgreSQL 的一个扩展，允许你在数据库中存储和查询向量 embeddings。让我们试试看。

首先，我们将启用 Vector 扩展。在 Supabase 中，可以通过 Database → Extensions 从网络门户完成此操作。你也可以通过运行以下 SQL 来完成：

create extension if not exists vector;

接下来，让我们创建一个表来存储我们的文档和它们的 embeddings：

create table documents (
  id bigserial primary key,
  content text,
  embedding vector(1536)
);

pgvector 引入了一个名为 vector 的新数据类型。在上面的代码中，我们创建了一个名为 embedding 的列，类型为 vector。向量的大小定义了向量保持多少个维度。OpenAI 的 text-embedding-ada-002 模型输出 1536 维，所以我们将使用它作为向量大小。

我们还创建了一个名为 content 的文本列来存储产生此 embedding 的原始文档文本。根据你的用例，你可能只需在这里存储文档的引用（URL 或外键）。

很快我们将需要对这些 embeddings 执行相似性搜索。让我们创建一个函数来完成这个任务：

create or replace function match_documents (
  query_embedding vector(1536),
  similarity_threshold float,
  match_count int
)
returns table (
  id bigint,
  content text,
  similarity float
) as $$
  select
    documents.id,
    documents.content,
    1 - (documents.embedding <=> query_embedding) as similarity
  from documents
  where 1 - (documents.embedding <=> query_embedding) > similarity_threshold
  order by documents.embedding <=> query_embedding
  limit match_count;
$$ language sql stable;

pgvector 引入了 3 个新的操作符来计算相似性：

<=>: 余弦距离

<#>: 负内积

@@: 内积

OpenAI 推荐在其 embeddings 上使用余弦相似性，所以我们在这里使用它。

现在我们可以调用 match_documents()，传入我们的 embedding、相似性阈值和匹配数，我们将获得所有匹配文档的列表。由于这一切都由 Postgres 管理，我们的应用代码变得非常简单。

一旦你的表开始增长 embeddings，你可能会想添加一个索引来加速查询。向量索引在排序结果时特别重要，因为向量不是按相似性分组的，所以通过顺序扫描找到最接近的向量是一个资源密集的操作。

每个距离操作符需要不同类型的索引。我们期望按余弦距离排序，所以我们需要 vector_cosine_ops 索引。一个好的起始列表数是 4 * sqrt(table_rows)：

create index on documents using ivfflat (embedding vector_cosine_ops)
  with (lists = 100);

你可以在 pgvector 的 GitHub 页面上阅读更多关于索引的信息。

生成 embeddings

让我们使用 JavaScript 来生成 embeddings 并将其存储在 Postgres 中：

import { createClient } from '@supabase/supabase-js'
import { Configuration, OpenAIApi } from 'openai'

const openai = new OpenAIApi(new Configuration({
  apiKey: process.env.OPENAI_API_KEY,
}))

const supabase = createClient(
  process.env.SUPABASE_URL,
  process.env.SUPABASE_KEY
)

const generateEmbeddings = async (content: string) => {
  const embeddingResponse = await openai.createEmbedding({
    model: 'text-embedding-ada-002',
    input: content,
  })

  const [{ embedding }] = embeddingResponse.data.data
  await supabase
    .from('documents')
    .insert({
      content,
      embedding,
    })
    .single()
}

generateEmbeddings('Hello World')

构建一个简单的搜索函数

最后，让我们创建一个 Edge Function 来执行我们的相似性搜索：

import { serve } from 'https://deno.land/std@0.131.0/http/server.ts'
import { createClient } from 'https://esm.sh/@supabase/supabase-js@2.1.1'
import { Configuration, OpenAIApi } from 'https://esm.sh/openai@3.1.0'

const openai = new OpenAIApi(new Configuration({
  apiKey: Deno.env.get('OPENAI_API_KEY'),
}))

const supabase = createClient(
  Deno.env.get('SUPABASE_URL'),
  Deno.env.get('SUPABASE_SERVICE_ROLE_KEY')
)

serve(async (req: Request) => {
  const { query } = await req.json()

  const embeddingResponse = await openai.createEmbedding({
    model: 'text-embedding-ada-002',
    input: query,
  })

  const [{ embedding }] = embeddingResponse.data.data

  const { data } = await supabase.rpc('match_documents', {
    query_embedding: embedding,
    similarity_threshold: 0.5,
    match_count: 10,
  })

  return new Response(JSON.stringify(data), {
    headers: { 'Content-Type': 'application/json' },
  })
})

构建一个更智能的搜索函数

ChatGPT 不仅仅返回现有文档。它能够将各种信息吸收到一个单一、连贯的答案中。为此，我们需要为 GPT 提供一些相关文档和一个 prompt，它可以用这些来制定答案。

OpenAI 的 text-davinci-003 completion 模型最大的挑战之一是 4000 token 的限制。你必须将你的 prompt 和生成的 completion 都控制在 4000 token 内。如果你想提示 GPT-3 回答关于你自己的自定义知识库的问题，而这些问题永远无法放入一个 prompt 中，这会很有挑战性。

Embeddings 可以通过将你的 prompts 分成两个阶段的过程来帮助解决这个问题：

查询你的 embedding 数据库，找到与问题相关的最相关文档

将这些文档作为上下文注入，供 GPT-3 在其答案中参考

这是另一个 Edge Function，它基于上面的简单例子进行了扩展：

import { serve } from 'https://deno.land/std@0.131.0/http/server.ts'
import { createClient } from 'https://esm.sh/@supabase/supabase-js@2.1.1'
import { Configuration, OpenAIApi } from 'https://esm.sh/openai@3.1.0'

const openai = new OpenAIApi(new Configuration({
  apiKey: Deno.env.get('OPENAI_API_KEY'),
}))

const supabase = createClient(
  Deno.env.get('SUPABASE_URL'),
  Deno.env.get('SUPABASE_SERVICE_ROLE_KEY')
)

serve(async (req: Request) => {
  const { query } = await req.json()

  // Step 1: Generate an embedding for the query itself
  const embeddingResponse = await openai.createEmbedding({
    model: 'text-embedding-ada-002',
    input: query,
  })

  const [{ embedding }] = embeddingResponse.data.data

  // Step 2: Retrieve the most relevant documents
  const { data: documents } = await supabase.rpc('match_documents', {
    query_embedding: embedding,
    similarity_threshold: 0.5,
    match_count: 10,
  })

  // Step 3: Inject the documents into the prompt
  const systemPrompt = `You are a helpful assistant. Answer the user's question based on the following documents:\n\n${documents
    .map((d) => d.content)
    .join('\n\n')}`

  // Step 4: Call the completion API
  const completionResponse = await openai.createChatCompletion({
    model: 'gpt-3.5-turbo',
    messages: [
      { role: 'system', content: systemPrompt },
      { role: 'user', content: query },
    ],
    max_tokens: 1024,
  })

  const [{ message }] = completionResponse.data.choices
  return new Response(JSON.stringify({ message }), {
    headers: { 'Content-Type': 'application/json' },
  })
})

OpenAI API 的响应时间取决于"答案"的长度。ChatGPT 通过立即将响应流式传输给用户来提供良好的用户体验。你可以在 Supabase 文档中看到类似的效果：

OpenAI API 支持使用 Server Side Events 的 completion 流式传输。Supabase Edge Functions 运行在 Deno 上，它也支持 Server Side Events。查看此提交以了解我们如何修改上述 Function 来构建流式接口。

在 Postgres 中存储 embeddings 开启了无限可能。你可以将搜索函数与遥测函数结合，添加用户提供的反馈（点赞/踩），并使你的搜索与产品更加紧密集成。

pgvector 扩展现已在所有新的 Supabase 项目上可用。要试用它，请启动一个新的 Postgres 数据库：database.new

更多 pgvector 和 AI 资源

Supabase Clippy: Supabase Docs 的 ChatGPT

Hugging Face 现已支持在 Supabase 中

如何使用 Supabase Edge Runtime 从头开始构建 ChatGPT 插件

Docs pgvector: Embeddings 和向量相似性

为 AI 工作负载选择计算附加组件

pgvector v0.5.0: 使用 HNSW 索引进行更快的语义搜索