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

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:@staticmethoddef 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_documentsrerank_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"] = scorererank_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_instancedef 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_idfor 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 documentrerank_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模型，可以根据需要进行选择。