Timescale/pgai项目：使用Vectorizer和LiteLLM评估嵌入模型的完整指南-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00635/article/details/148508517

Timescale/pgai项目：使用Vectorizer和LiteLLM评估嵌入模型的完整指南

pgai Helper functions for AI workflows 项目地址: https://gitcode.com/gh_mirrors/pg/pgai

概述

在现代AI应用中，嵌入模型（Embedding Models）扮演着至关重要的角色，它们能够将文本转换为向量表示，从而支持语义搜索、推荐系统等多种应用场景。本文将详细介绍如何在Timescale/pgai项目中，利用其Vectorizer功能和LiteLLM集成来评估不同嵌入模型的性能表现。

技术背景

pgai是PostgreSQL的一个扩展，提供了强大的AI功能集成。其中，Vectorizer是一个关键组件，它能够：

自动处理文本数据的分块（chunking）
调用各种嵌入模型生成向量表示
将结果存储在数据库表中
支持多种模型的无缝切换

LiteLLM则是一个统一的API接口，允许开发者通过单一接口访问多种大语言模型服务，包括Cohere、Mistral、OpenAI等。

评估方案设计

本评估方案采用了Paul Graham的散文集作为测试数据集，通过生成多种类型的问题来全面测试不同嵌入模型的检索能力。评估的核心思路是：

从文本中随机选取多个片段（chunks）
为每个片段生成不同类型的问题
使用各嵌入模型检索最相关的文本片段
比较检索结果与原始片段的匹配程度

准备工作

环境配置

确保已安装Docker和Docker Compose
准备PostgreSQL数据库并安装pgai扩展
获取Paul Graham散文集数据集
准备各模型API密钥：
- Cohere API密钥
- Mistral API密钥
- OpenAI API密钥
- HuggingFace API密钥

关键参数配置

NUM_CHUNKS = 20               # 评估的随机文本块数量
NUM_QUESTIONS_PER_CHUNK = 20  # 每个文本块生成的问题总数
TOP_K = 10                    # 检索的最相似文本块数量
QUESTION_DISTRIBUTION = {     # 问题类型分布
    'short': 4,    # 简短问题（少于10个词）
    'long': 4,     # 需要详细回答的长问题
    'direct': 4,   # 关于明确信息的问题
    'implied': 4,  # 需要上下文理解的问题
    'unclear': 4   # 模糊或模棱两可的问题
}

实施步骤详解

1. 数据库环境搭建

使用提供的Docker Compose文件快速搭建环境，确保在配置文件中正确设置所有API密钥。

2. 初始化pgai扩展

连接到数据库后，执行以下SQL命令启用pgai扩展：

CREATE EXTENSION IF NOT EXISTS ai CASCADE;

3. 加载数据集

将Paul Graham的散文集加载到数据库中：

SELECT ai.load_dataset('sgoel9/paul_graham_essays');

4. 创建Vectorizer

为每个嵌入模型创建Vectorizer：

Cohere模型

SELECT ai.create_vectorizer(
    'paul_graham_essays'::regclass,
    destination => ai.destination_table('essays_cohere_embeddings'),
    embedding => ai.embedding_litellm(
        'cohere/embed-english-v3.0',
        1024,
        api_key_name => 'COHERE_API_KEY'
    ),
    chunking => ai.chunking_recursive_character_text_splitter('text', 512, 50)
);

Mistral模型

SELECT ai.create_vectorizer(
    'paul_graham_essays'::regclass,
    destination => ai.destination_table('essays_mistral_embeddings'),
    embedding => ai.embedding_litellm(
        'mistral/mistral-embed',
        1024,
        api_key_name => 'MISTRAL_API_KEY'
    ),
    chunking => ai.chunking_recursive_character_text_splitter('text', 512, 50)
);

OpenAI模型

SELECT ai.create_vectorizer(
    'paul_graham_essays'::regclass,
    destination => ai.destination_table('essays_openai_small_embeddings'),
    embedding => ai.embedding_openai(
        'text-embedding-3-small', 
        1024, 
        api_key_name => 'OPENAI_API_KEY'
    ),
    chunking => ai.chunking_recursive_character_text_splitter('text', 512, 50)
);

5. 监控处理状态

使用以下命令检查Vectorizer处理状态：

SELECT * FROM ai.vectorizer_status;

查看可能的错误信息：

SELECT * FROM ai.vectorizer_errors;

评估执行流程

随机选择文本块：从已向量化的表中随机选取指定数量的文本块
生成评估问题：使用GPT-4o-mini为每个文本块生成不同类型的问题
模型评估：对每个嵌入模型执行以下操作：
- 将问题转换为向量
- 在向量空间中搜索最相似的文本块
- 记录检索结果与原始文本块的匹配程度

评估结果分析

评估完成后，系统会生成多个输出文件：

chunks.csv：包含随机选择的原始文本块
questions.csv：为每个文本块生成的问题列表
results.csv：各模型的整体性能指标
detailed_results.csv：每个问题的详细评估结果

通过这些数据，可以分析：

不同模型在各种问题类型上的表现差异
模型对明确信息与隐含信息的理解能力
模型处理模糊问题的能力
各模型在检索准确率和召回率上的表现

技术要点解析

文本分块策略：使用递归字符文本分割器（recursive character text splitter）将长文本分割为512字符的块，重叠50字符
向量维度：所有模型统一使用1024维向量，确保比较的公平性
问题多样性：通过设计不同类型的问题，全面评估模型能力：
- 简短问题测试模型的精确匹配能力
- 长问题测试模型的上下文理解能力
- 直接问题测试事实检索能力
- 隐含问题测试语义理解能力
- 模糊问题测试模型的鲁棒性