基于 PostgreSQL pgvector 的多模态搜索

hudson2022
9 阅读7分钟

基于 PostgreSQL pgvector 的多模态搜索

摘要: 本文演示如何使用 PostgreSQL 和 pgvector 构建多模态搜索应用，允许用户使用图片或文本输入搜索印度食谱数据库。该应用利用多模态大语言模型从视觉和文本数据生成嵌入向量，实现通过存储在 PostgreSQL 中的向量嵌入进行相似性搜索。

原文链接

大型语言模型（LLM）已经显著超越了生成文本提示回复的能力。这些模型现已被训练具备高级能力，如解读图像和根据视觉输入提供详细描述。这为用户提供了更大的搜索能力。

在本文中，我将演示如何构建具有多模态搜索功能的应用。该应用的用户可以上传图片或提供文本输入来搜索印度菜谱数据库。应用构建为支持多种 LLM 提供商，允许用户在 OpenAI 或本地运行的模型之间进行选择（使用 Ollama）。文本嵌入随后使用 pgvector 存储在 PostgreSQL 中并进行查询。

要查看完整的源代码以及构建和运行应用的说明，请访问 GitHub 上的示例应用。

应用的完整演示和架构说明也可在 YouTube 上找到：

image.png

构建模块

在深入代码之前，让我们概述每个组件在构建多模态搜索应用中的作用。

image.png

  • 多模态大语言模型 (LLM)：在大型数据集上训练的模型，能够处理多种类型的数据，如文本、图像和语音。
  • 嵌入模型：将输入转换为固定维度数值向量的模型，用于相似性搜索。例如，OpenAI 的 text-embedding-3-small 模型生成 1536 维向量。
  • PostgreSQL：适用于各种应用的开源通用关系型数据库，配备了用于在 AI 应用中存储和查询向量嵌入的扩展。
  • pgvector：用于处理向量相似性搜索的 PostgreSQL 扩展。

现在我们已经了解了应用架构和基础组件，让我们将它们组合在一起！

生成和存储嵌入向量

本项目提供了从您选择的提供商生成嵌入向量的实用函数。让我们逐步介绍生成和存储文本嵌入向量的步骤。

读取包含原始数据集的 cuisines.csv 文件并将其存储在 Pandas DataFrame 中以便进行处理。

每个食谱的描述被传递给 generate_embedding 函数，以在 DataFrame 中填充新列 embeddings。然后将数据写入新的 output.csv 文件，包含用于相似性搜索的嵌入向量。

稍后，我们将更详细地了解 generate_embedding 函数的工作原理。

import sys
import os
import pandas as pd

# Add the project root to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from backend.llm_interface import generate_embedding

# Load the CSV file
csv_path = './database/cuisines.csv'
df = pd.read_csv(csv_path)

# Generate embeddings for each description in the CSV
df['embeddings'] = df['description'].apply(generate_embedding, args=(True,))

# Save the DataFrame with embeddings to a new CSV file
output_csv_path = './database/output.csv'
df.to_csv(output_csv_path, index=False)

print(f"Embeddings generated and saved to {output_csv_path}")

使用 pgvector，这些嵌入向量可以轻松存储在 PostgreSQL 的 embeddings 列中，类型为 vector。

CREATE EXTENSION IF NOT EXISTS vector;

CREATE table recipes (
   id SERIAL PRIMARY KEY,
   name text,
   description text,
   ...
   embeddings vector (768)
);

生成的 output.csv 文件可以使用 COPY 命令复制到数据库，也可以使用 Pandas DataFrame 提供的 to_sql 函数。

# Copy to recipes table running in Docker
docker exec -it postgres bin/psql -U postgres -c "\COPY recipes(name,description,...,embeddings) from '/home/database/output.csv' DELIMITER ',' CSV HEADER;"

# Write the DataFrame to the recipes table table
engine = create_engine('postgresql+psycopg2://username:password@hostname/postgres')

df.to_sql('recipes', engine, if_exists='replace', index=False)

有了存储食谱描述向量嵌入的 PostgreSQL 实例，我们就可以运行应用并执行查询了。

多模态搜索应用

让我们将应用连接到数据库，开始对食谱描述嵌入向量执行查询。

搜索端点接受通过多部分表单提交的文本和图像。

# server.py
from llm_interface import describe_image, generate_embedding

...

@app.route('/api/search', methods=['POST'])
def search():
   image_description = None
   query = None
   # multipart form data payload
   if 'image' in request.files:
       image_file = request.files['image']
       image_description = describe_image(image_file)

   data = request.form.get('data')
   if data and 'query' in data:
       try:
           data = json.loads(data)
           query = data['query']
       except ValueError:
           return jsonify({'error': 'Invalid JSON data'}), 400

   if not image_description and not query:
       return jsonify({'error': 'No search query or image provided'}), 400

   embedding_query = (query or '') + " " + (image_description or '')

   embedding = generate_embedding(embedding_query)

   try:
       conn = get_db_connection()
       cursor = conn.cursor()
       cursor.execute("SELECT id, name, description, instructions, image_url FROM recipes ORDER BY embeddings <=> %s::vector  LIMIT 10", (embedding,))
       results = cursor.fetchall()
       cursor.close()
       conn.close()

       return jsonify({'results': results, 'image_description': image_description or None})

   except Exception as e:
       return jsonify({'error': str(e)}), 500

虽然这个 API 非常直接，但有两个值得关注的辅助函数：describe_image 和 generate_embedding。让我们更详细地了解这些函数的工作原理。

# llm_interface.py
# Function to generate a description from an image file
def describe_image(file_path):
   image_b64 = b64_encode_image(file_path)
   custom_prompt = """You are an expert in identifying Indian cuisines.
   Describe the most likely ingredients in the food pictured, taking into account the colors identified.
   Only provide ingredients and adjectives to describe the food, including a guess as to the name of the dish.
   Output this as a single paragraph of 2-3 sentences."""


   if(LLM_ECOSYSTEM == 'ollama'):
       response = ollama.generate(model=LLM_MULTIMODAL_MODEL, prompt=custom_prompt, images=[image_b64])
       return response['response']
   elif(LLM_ECOSYSTEM == 'openai'):    
       response = client.chat.completions.create(messages=[
           {"role": "system", "content": custom_prompt},
           {"role": "user", "content": [
           {
               "type": "image_url",
               "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
           }]}
       ], model=LLM_MULTIMODAL_MODEL)
       return response.choices[0].message.content
  else:
       return "No Model Provided"

describe_image 函数接受图像文件路径，并将 base64 编码发送给用户首选的 LLM。

为简单起见，该应用目前支持在 Ollama 本地运行的模型，或通过 OpenAI 可用的模型。这个 base64 图像表示伴随着一个自定义提示，告诉 LLM 充当印度菜专家，以便准确描述上传的图像。在使用 LLM 时，清晰的提示构建对于获得期望的结果至关重要。

从函数返回的图像简短描述可以传递给 generate_embedding 函数，生成存储在数据库中的向量表示。

# llm_interface.py
# Function to generate embeddings for a given text
def generate_embedding(text):
   if LLM_ECOSYSTEM == 'ollama':
       embedding = ollama.embeddings(model=LLM_EMBEDDING_MODEL, prompt=text)
       return embedding['embedding']
   elif LLM_ECOSYSTEM == 'openai':
       response = client.embeddings.create(model=LLM_EMBEDDING_MODEL, input=text)
       embedding = response.data[0].embedding
       return embedding
   else:
       return "No Model Provided"

generate_embedding 函数依赖于 AI 生态系统中的另一类模型，这些模型从文本生成向量嵌入。这些模型也可以通过 Ollama 和 OpenAI 获得，分别返回 768 和 1536 维。

通过为 LLM 返回的每个图像描述生成嵌入向量（以及通过表单输入可选地提供额外文本），API 端点可以使用 pgvector 中的余弦距离进行查询，以提供准确的结果。

cursor.execute("SELECT id, name, description, instructions, image_url FROM recipes ORDER BY embeddings <=> %s::vector  LIMIT 10", (embedding,))
results = cursor.fetchall()

通过连接 UI 并通过图像和简短文本描述进行搜索，该应用可以利用 pgvector 在数据集上执行相似性搜索。

分布式 SQL 在 AI 应用中的优势

让我们探讨如何利用分布式 SQL 使我们的应用更具可扩展性和弹性。

以下是使用 pgvector 的 AI 应用受益于分布式 PostgreSQL 数据库（如 YugabyteDB）的几个关键原因：

  1. 嵌入向量消耗大量存储和内存。 1536 维的 OpenAI 模型存储 1000 万条记录需要约 57GB 的空间。水平扩展提供了存储向量所需的空间。
  2. 向量相似性搜索是计算密集型的。 通过扩展到多个节点，应用可以访问无限制的 CPU 和 GPU 资源。
  3. 避免服务中断。 数据库能够抵御节点、数据中心和区域故障，因此 AI 应用永远不会因数据库层而经历停机。

YugabyteDB 是构建在 PostgreSQL 之上的分布式 SQL 数据库。它与 PostgreSQL 在功能和运行时上兼容，允许您重复使用为标准版 PostgreSQL 创建的库、驱动程序、工具和框架。

YugabyteDB 具有 pgvector 兼容性，并提供原生 PostgreSQL 中的所有功能。这使其非常适合希望提升 AI 应用能力的人。

结论

使用 AI 生态系统中最新的多模态模型，使应用添加图像搜索变得轻而易举。这个简单但强大的应用展示了 PostgreSQL 支持的应用如何轻松支持最新、最强大的 AI 功能。

想了解更多关于 AI 的信息吗？请查看以下最新博客：

avatar
文章
阅读
粉丝