大模型 RAG：基於 PgSql 的向量檢索

一 RAG 與向量檢索

1.1 RAG 概念

檢索增強生成（Retrieval-augmented Generation），簡稱 RAG。

RAG 通常包括兩個階段：1、檢索上下文相關信息；2、使用檢索到的知識指導生成過程。簡單來說，就像開卷考試，我們可以攜帶參考材料用來查找相關信息來回答問題。

1.2 RAG 意義

目前對大模型的使用通常存在兩個主要挑戰：1、由於生成模型依賴於內在知識（權重），對於未覆蓋到的知識領域可能會產生大量的幻覺，也就是 “一本正經的胡說八道”；其次，由於大模型參數量巨大，重新訓練或微調方法成本過高。

1.3 RAG 應用框架

RAG 應用框架如下圖所示，包含文本向量化、檢索向量數據庫、獲取上下文（相關知識 / 文檔）、prompt 構造、調用大模型執行文本生成等主要環節。

當然這裏只是描述了一個主流程，實際應用還有對多輪 / 歷史會話處理、多輪問題改寫，意圖識別與分發，結果聚合等等。

二 向量數據庫

    由於介紹 RAG 的文章已經很多，本篇主要打算介紹向量數據庫部分，所以不再對 RAG 進行贅述。

2.1 向量數據庫方案

目前已經有十幾種可選的向量數據庫技術方案，包括 Milvus、MongoDB Atlas、Chroma、Weaviate 等。從分類的角度說，包括：原生向量數據庫（Chroma、LanceDB、Mivus 等）、支持向量的全文檢索數據庫（Elastic、Lucene、OpenSearch 和 Solr）、支持向量的 NoSQL 數據庫（Cassandra、Rockset、Azure Cosmos DB 和 MongoDB）、和支持向量的關係數據庫（PostgreSQL、Clickhouse、SingleStoreDB 等）。

主流向量數據庫的部分指標對比如下：

通常大模型應用場景，Milvus 等原生向量數據庫是最推薦的。但根據實際使用場景和習慣，在知識庫數量並不是很大的場景（億級以下），考慮到使用習慣和學習複雜度，也可以選擇 PgSQL 等關係型數據庫。

2.2 PgSql 與 PgVector

Postgres 通過_ pg_vector 和 pg_embdding_ 兩個插件來實現向量數據庫，讓 PG 數據庫支持向量索引檢索的能力。其索引算法使用的是基於 Faiss 的 IVF Flat 索引，提供了優異的召回率。

三 基於 PgSql 的向量檢索示例

3.1 建立向量庫

如下建表語句所示，向量庫主要文檔內容和 embedding（文檔向量化結果）。考慮到在實際應用場景可能涉及文檔權限控制，或展示參考資料，我們還設計了文檔名、文檔 id 字段。

-- vector.knowledge_doc_vector definition
-- Drop table
-- DROP TABLE vector.knowledge_doc_vector;
CREATE TABLE vector.knowledge_doc_vector (
  id bigserial NOT NULL,
  doc_id bigserial NOT NULL,
  embedding public.vector NULL,
  doc_content text NULL,
  doc_name varchar NULL,
  doc_page varchar NULL,
  CONSTRAINT knowledge_doc_vector_pkey PRIMARY KEY (id)
);

示例數據如下：

3.2 Mybatis 連接 PgSql

除了 mybatis-plus 之外，引入 pgsql 和 pgvector 兩個依賴：

<dependency>
            <groupId>com.baomidou</groupId>
            <artifactId>mybatis-plus-boot-starter</artifactId>
            <version>3.5.3.1</version>
        </dependency>
     <dependency>
            <groupId>org.postgresql</groupId>
            <artifactId>postgresql</artifactId>
            <version>42.6.0</version>
        </dependency>
        <dependency>
            <groupId>com.pgvector</groupId>
            <artifactId>pgvector</artifactId>
            <version>0.1.6</version>
        </dependency>

3.3 編寫 mapper

用於查找相似向量，按照餘弦相似度計算

import com.baomidou.mybatisplus.core.mapper.BaseMapper;
import com.fawvw.pms.vector.domain.dto.KnowledgeDocVectorDTO;
import com.fawvw.pms.vector.domain.entity.KnowledgeDocVector;
import com.google.gson.Gson;
import org.apache.ibatis.annotations.Mapper;
import org.apache.ibatis.annotations.SelectProvider;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
@Mapper
public interface VectorMapper extends BaseMapper<KnowledgeDocVector> {
    @SelectProvider(type = SqlProvider.class, method = "findUsersWithSimilarVectors")
    List<KnowledgeDocVectorDTO> findUsersWithSimilarVectors(double[] targetVector);
    static class SqlProvider {
        public String findUsersWithSimilarVectors(Map<String, Object> param) {
            System.out.println("param keys:" + new Gson().toJson(param.keySet()));
            double[] targetVector = (double[]) param.get("array");
            String sql = "SELECT id, doc_id, embedding, doc_content, " +
                    "cosine_distance(embedding, '" + Arrays.toString(targetVector) + "') AS similarity " +
                    "FROM knowledge_doc_vector " +
                    "ORDER BY similarity DESC";
            return sql;
        }
    }
}

3.4 方法中調用，查詢相似向量

下面示例代碼使用的是阿里的靈積服務計算 embedding，也可以採用其他方式計算。

public List findSimilarVector(String query) {
        // query轉爲向量
        TextEmbeddingResult embeddingResult = dashCodeHttpApi.textEmbedding(query);
        List<Double> embeddings = embeddingResult.getOutput().getEmbeddings().get(0).getEmbedding();
        double[] vector = new double[embeddings.size()];
        for (int i = 0; i < embeddings.size(); i++) {
            Double embedding = embeddings.get(i);
            vector[i] = embedding;
        }
        log.info("vector:{}", vector);
        List<KnowledgeDocVectorDTO> vectorDTOList = vectorMapper.findUsersWithSimilarVectors(vector);
        log.info("vectorDTOList:{}", vectorDTOList);
        // 過濾掉相似度低的向量檢索結果 閾值:0.7
        vectorDTOList = vectorDTOList.stream()
                .filter(vectorDTO -> vectorDTO.getSimilarity() <= 0.7)
                .collect(Collectors.toList());
        return vectorDTOList;
    }

拿到與問題相關的文檔知識之後，我們就可以封裝 prompt，並調用大模型 API 獲取生成式回答了。具體方法留在下一篇中給出。

四 小結

    本篇介紹了 RAG 的一些基礎知識，以及向量庫在其中的作用及選型。並基於 PgSQL 給出了一個向量庫的使用示例。接下來的文章中，我們將深入探索如何實現可用的 RAG 應用。歡迎留言一起探討。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/SIkuvKUWW_NmZRPGbgCIkA

猜你喜歡