来自工业界的知识库 RAG 服务(五)，模块化知识库 GoMate 实现方案详解

news2024/10/6 8:23:00

背景介绍

很早以前就看过一篇介绍 RAG 的综述性文章 Retrieval-Augmented Generation for Large Language Models: A Survey, 其中介绍了 RAG 的模块化架构：
请添加图片描述

但是一直没有看到对应的实现方案，以前也没有理解此框架的优势之处。随着在相关领域的持续实践，逐渐理解到模块化 RAG 的优势。模块化 RAG 可以方便地进行服务的迭代升级，另外模块可以自由组合，灵活性更好。

最近注意到来自中科院的 RAG 服务 GoMate 开源了，核心亮点就是模块化 RAG 设计，本文就深入介绍来 GoMate 的框架设计与实现细节。

框架设计

之前介绍的 RAG 框架存在着不同的侧重点，来自有道的 QAnything 强调其 Rerank 设计，RagFlow 强调其精细的文件处理。而 GoMate 则主要强调其模块化设计。

GoMate 的框架设计图如下所示：

请添加图片描述

可以看到，框架中将 RAG 流程拆分为不同模块，熟悉 RAG 服务还是比较容易理解的，一些核心模块如下所示：

Retrieval 对应的就是知识库的检索
Rerank&Judge 对应的就是检索内容的重排序与过滤
Generation 对应的就是大模型内容生成
Citation 对应的就是检索内容引用的处理

每个模块可以存在多种不同的实现形式，比如 Retrieval 图中包含如下所示的实现：

Sparse Retriever 对应的是稀疏检索，目前常见的 BM25 的检索；
Dense Retriever 对应的密集检索，RAG 中最常见的是向量检索；
Ensemble Retriever 对应的应该是混合检索；

从架构图可以看到 GoMate 就是包含了多个基础的功能模块，通过功能模块的的组合实现具体的应用。

实现细节

GoMate 是从 2024 年 2 月开始开发的，截止目前经过了 4 个月的开发，虽然功能还不够完善，但是已经可以看到模块化的完整设计思想了。

项目主要包含两部分:

modules 包含的就是基础模块, 与上面架构图中的模块可以对应起来
applications 表示的是模块组合生成的应用，实现业务所需的功能

功能模块

每个模块对应的职责比较明确，以 RAG 中核心的检索模块为例，深入查看对应的实现：

基础的检索接口在 gomate/modules/retrieval/base.py 中定义，只包含一个 retrieve() 方法：

class BaseRetriever(ABC):
    @abstractmethod
    def retrieve(self, query: str) -> str:
        pass

BM25 检索算法在 gomate/modules/retrieval/bm25_retriever.py 实现，其中实现了 BM25Okapi, BM25L 和 BM25Plus 版本的算法，具体如下所示：

class BM25Retriever(BaseRetriever):
    # 实现不同的 BM25 算法

    def build_from_texts(self, corpus):
        self.corpus=corpus
        if self.algorithm == 'Okapi':
            self.bm25 = BM25Okapi(corpus=corpus, tokenizer=self.tokenizer, k1=self.k1, b=self.b, epsilon=self.epsilon)
        elif self.algorithm == 'BM25L':
            self.bm25 = BM25L(corpus=corpus, tokenizer=self.tokenizer, k1=self.k1, b=self.b, delta=self.delta)
        elif self.algorithm == 'BM25Plus':
            self.bm25 = BM25Plus(corpus=corpus, tokenizer=self.tokenizer, k1=self.k1, b=self.b, delta=self.delta)
        else:
            raise ValueError('Algorithm not supported')

    # 检索接口的实现

    def retrieve(self, query: str='',top_k:int=3) -> List[Dict]:
        tokenized_query = " ".join(self.tokenizer(query))
        search_docs = self.bm25.get_top_n(tokenized_query, self.corpus, n=top_k)
        return search_docs

向量检索算法是基于 Faiss 实现的，简化后实现如下所示：

class DenseRetriever(BaseRetriever):
    # 初始化 Faiss 向量检索

    def __init__(self, config):
        self.config = config
        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)
        self.model = AutoModel.from_pretrained(config.model_name)
        self.index = faiss.IndexFlatIP(config.dim)

    # 检索接口实现

    def retrieve(self, query):
        D, I = self.index.search(self.get_embedding([query]), self.top_k)
        return [{self.doc_map[idx]: score} for idx, score in zip(I[0], D[0]) if idx in self.doc_map]

在 gomate/modules/retrieval/tree_retriever.py 中还实现了 RAPTOR 检索，感兴趣的也可以深入查看实现细节。

可以看到检索模块包含的就是不同算法的实现，上层应用可以根据需要选择合适的检索算法。

其他的模块与检索模块类似，对应的就是单个功能的实现。

应用

目前实现的应用比较少，只有 Rerank，Rewrite 和 RAG ，目前只有 RAG 属于一个完整的应用，其他都是单个模块的测试，深入查看 RAG 应用的实现，具体如下所示：

class RagApplication():
    # 初始化向量库

    def init_vector_store(self):
        docs=self.reader.get_content(max_token_len=600, cover_content=150)
        self.vector_store.document=docs
        self.vector_store.get_vector(EmbeddingModel=self.embedding_model)
        self.vector_store.persist(path='storage')
        self.vector_store.load_vector(path='storage')

    # 向量库中添加文件

    def add_document(self, file_path):
        docs = self.reader.get_content_by_file(file=file_path, max_token_len=512, cover_content=60)
        self.vector_store.add_documents(self.config.vector_store_path, docs, self.embedding_model)

    # RAG 对话

    def chat(self, question: str = '', topk: int = 5):
        # 向量检索

        contents = self.vector_store.query(question, EmbeddingModel=self.embedding_model, k=topk)
        content = '\n'.join(contents[:5])

        # 大模型聊天

        response, history = self.llm.chat(question, [], content)
        return response, history,contents