dify实现分析-rag-内容检索rerank的两种实现

概述

本文介绍dify的rerank的实现原理。rerank可以对检索到的多个文档进行重排序。通过重排序可以找到更加合适的文档分块。

dify中的两种rerank实现方式

dify中提供两种rerank的方式：

• 使用rerank模型：一种是通过rerank模型来直接对输入文档列表(documents)和问题(query)的相关性进行评估，打分，并返回与问题最合适的文档列表。
• 根据权重计算：嵌入向量以及文档分词与query相似度的综合分数。先对文档和query进行分词处理，然后计算每个文档和query的相似度分数；在对文档列表以及query进行向量化处理，并计算每个文档和query向量的相似度分数。然后根据权重把两者的分数和权重相乘，并加起来。就得到了每个文档的最终相似度。然后取分数最大的top_n个文档。

rerank的两种rerank模式的定义代码如下：

class RerankRunnerFactory:
    @staticmethod
    def create_rerank_runner(runner_type: str, *args, **kwargs) -> BaseRerankRunner:
        match runner_type:
            case RerankMode.RERANKING_MODEL.value:
                return RerankModelRunner(*args, **kwargs)
            case RerankMode.WEIGHTED_SCORE.value:
                return WeightRerankRunner(*args, **kwargs)
            case _:
                raise ValueError(f"Unknown runner type: {runner_type}")

权重rerank的实现:WeightRerankRunner

基于权重的rerank是在WeightRerankRunner.run()函数中实现的，该函数的实现逻辑如下：

1. 相对document进行去重处理
2. 使用BM25算法对query和每个doucment进行分词处理，并计算每个文档和query的相似度分数。这样就得到每个document和query的相似度分数。
3. 对query和document进行向量化处理，并计算每个document和query向量的相似度分数。这样就得到每个document和query的向量的相似度分数。
4. 针对每个document：将基于分词和向量的两个分数按照给定的权重综合计算，得到每个document和query相似度的最终评分。
5. 根据document的得分对其进行排序，并选择得分最高的top_n的个document返回。

WeightRerankRunner中rerank的实现代码如下：

class WeightRerankRunner(BaseRerankRunner):
	def run(
        self,
        query: str,
        documents: list[Document],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
        user: Optional[str] = None,
    ) -> list[Document]:
        """
        Run rerank model
        :param query: search query
        :param documents: documents for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :param user: unique user id if needed

        :return:
        """
        docs = []
        doc_id = []
        unique_documents = []
        # 对文档列表去重
        for document in documents:
            if document.metadata["doc_id"] not in doc_id:
                doc_id.append(document.metadata["doc_id"])
                docs.append(document.page_content)
                unique_documents.append(document)

        documents = unique_documents

        rerank_documents = []
        # 对query和每个文档分块进行分词处理，并计算每个文档和query的相似度分数。
        query_scores = self._calculate_keyword_score(query, documents)

        # 对query和document进行向量化处理，并计算每个document和query向量的相似度分数
        query_vector_scores = self._calculate_cosine(self.tenant_id, query, documents, self.weights.vector_setting)
        # 针对每个document：将基于分词和向量的两个分数按照给定的权重综合计算，得到最终的评分
        for document, query_score, query_vector_score in zip(documents, query_scores, query_vector_scores):
            # format document
            # 计算综合得分
            score = (
                self.weights.vector_setting.vector_weight * query_vector_score
                + self.weights.keyword_setting.keyword_weight * query_score
            )
            # 如果设置了 score_threshold 且文档的综合评分为负数，则跳过
            if score_threshold and score < score_threshold:
                continue
            # 添加文档的综合评分到元数据中
            document.metadata["score"] = score
            rerank_documents.append(document)
        # 将所有符合要求的文档按分数从高到低排序。
        rerank_documents = sorted(rerank_documents, key=lambda x: x.metadata["score"], reverse=True)
        # 根据 top_n 参数返回评分最高的前 N 个文档；如果不设置 top_n，则返回所有符合条件的文档。
        return rerank_documents[:top_n] if top_n else rerank_documents

使用rerank模型的实现：RerankModelRunner

使用rerank模型来实现document的rerank的实现比较简单，实现逻辑如下：

1. 对document进行去重处理，若是dify文档，记录doc_id。
2. 调用rerank模型来对document进行重拍处理，并输出重排序后的document列表。这里要注意，rerank模型的输入参数：

     rerank_result = self.rerank_model_instance.invoke_rerank(
            query=query, docs=docs, score_threshold=score_threshold, top_n=top_n, user=user
        )

3. 格式化重新排序后的document，生成Document对象，并把排序后的分数添加到document.metadata[“score”]字段中。

说明：rerank模型是用户通过页面配置的。

使用rerank模型实现的代码如下：

class RerankModelRunner(BaseRerankRunner):
    def __init__(self, rerank_model_instance: ModelInstance) -> None:
        self.rerank_model_instance = rerank_model_instance

    def run(
        self,
        query: str,
        documents: list[Document],
        score_threshold: Optional[float] = None,
        top_n: Optional[int] = None,
        user: Optional[str] = None,
    ) -> list[Document]:
        """
        Run rerank model
        :param query: search query
        :param documents: documents for reranking
        :param score_threshold: score threshold
        :param top_n: top n
        :param user: unique user id if needed
        :return:
        """
        docs = []
        doc_id = set()
        unique_documents = []
        # 对文档列表去重，针对dify的文档,记录doc_id
        for document in documents:
            if document.provider == "dify" and document.metadata["doc_id"] not in doc_id:
                doc_id.add(document.metadata["doc_id"])
                docs.append(document.page_content)
                unique_documents.append(document)
            elif document.provider == "external":
                if document not in unique_documents:
                    docs.append(document.page_content)
                    unique_documents.append(document)

        documents = unique_documents

        # 调用rerank模型对文档进行重排，该方法返回一个包含重新排序后的document（如文本和索引）
        rerank_result = self.rerank_model_instance.invoke_rerank(
            query=query, docs=docs, score_threshold=score_threshold, top_n=top_n, user=user
        )

        rerank_documents = []

        # 格式化重新排序后的文档
        for result in rerank_result.docs:
            # format document
            rerank_document = Document(
                page_content=result.text,
                metadata=documents[result.index].metadata,
                provider=documents[result.index].provider,
            )
            # 将重新排序后的分数添加到新文档的元数据中
            rerank_document.metadata["score"] = result.score
            # 添加到rerank_documents结果列表中
            rerank_documents.append(rerank_document)

        return rerank_documents

说明：dify支持多种rerank模型的实现，都继承了RerankModel类。

总结

本文介绍了dify的两种rerank方式的实现。目前dify中支持2种rerank方式，一种通过rerank模型来实现，一种是通过计算分词和向量与query的相似度来实现。另外，dify支持多种rerank模型，可以根据需要进行选择。