RAG Fusion 的简单尝试

0. RAG Fusion 是什么

RAG-Fusion是一个旨在弥合传统搜索模式与人类查询的多面维度之间差距的搜索方法。它受到增强检索生成(RAG)能力的启发,通过采用多种查询生成和互补排序融合来重新排列搜索结果,使之更进一步。

用通俗的语言来说，就是让大语言模型针对我们的问题，在生成几个类似问题，然后使用原始问题&大语言模型生成的问题去做相似性检索，最后在对所有的检索结果进行打分和排序。

1. 示例代码

import os
import openai
import random

# 初始化 OpenAI API
os.environ["OPENAI_API_KEY"] = "EMPTY"
os.environ["OPENAI_API_BASE"] = "http://192.168.31.15:8000/v1"
openai.api_key = "EMPTY"
openai.api_base = "http://192.168.31.15:8000/v1"

if openai.api_key is None:
    raise Exception("没有找到 OpenAI API 密钥。请将其设置为环境变量或在 main.py 中设置。")


# 使用OpenAI的ChatGPT生成查询的功能
def generate_queries_chatgpt(original_query):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "system",
             "content": "You are a helpful assistant that generates multiple search queries based on a single input query."},
            {"role": "user", "content": f"Generate multiple search queries related to: {original_query}"},
            {"role": "user", "content": "OUTPUT (4 queries):"}
        ]
    )

    generated_queries = response.choices[0]["message"]["content"].strip().split("\n")
    print("生成的查询：")
    for query in generated_queries:
        print(query)

    return generated_queries


# Mock function to simulate vector search, returning random scores
def vector_search(query, all_documents):
    available_docs = list(all_documents.keys())
    random.shuffle(available_docs)
    selected_docs = available_docs[:random.randint(2, 5)]
    scores = {doc: round(random.uniform(0.7, 0.9), 2) for doc in selected_docs}
    return {doc: score for doc, score in sorted(scores.items(), key=lambda x: x[1], reverse=True)}


# Reciprocal Rank Fusion algorithm
def reciprocal_rank_fusion(search_results_dict, k=60):
    fused_scores = {}
    print("初始各个查询搜索结果排名：")
    for query, doc_scores in search_results_dict.items():
        print(f"查询： '{query}': {doc_scores}")

    for query, doc_scores in search_results_dict.items():
        for rank, (doc, score) in enumerate(sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)):
            if doc not in fused_scores:
                fused_scores[doc] = 0
            previous_score = fused_scores[doc]
            fused_scores[doc] += 1 / (rank + k)
            print(f"根据查询 '{query}' 中的排名更新 {doc} 的得分，从 {previous_score} 更新为 {fused_scores[doc]} ")

    reranked_results = {doc: score for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)}
    print("最终重新排序的结果：", reranked_results)
    return reranked_results


# Dummy function to simulate generative output
def generate_output(reranked_results, queries):
    return f"根据查询 {queries} 和重新排序的文档的最终输出：{list(reranked_results.keys())}"


# Predefined set of documents (usually these would be from your search database)
all_documents = {
    "doc1": "气候变化与经济影响.",
    "doc2": "由于气候变化导致的公共卫生问题.",
    "doc3": "气候变化：一种社会角度.",
    "doc4": "应对气候变化的技术解决方案.",
    "doc5": "必需的政策变化来抗击气候变化.",
    "doc6": "气候变化对生物多样性的影响.",
    "doc7": "气候变化：科学与模型.",
    "doc8": "全球蜿蜒：气候变化的子集.",
    "doc9": "气候变化如何影响日常天气.",
    "doc10": "气候变化行动史."
}

# Main function
if __name__ == "__main__":
    original_query = "气候变化的影响"
    generated_queries = generate_queries_chatgpt(original_query)

    all_results = {}
    for query in generated_queries:
        search_results = vector_search(query, all_documents)
        all_results[query] = search_results

    reranked_results = reciprocal_rank_fusion(all_results)

    final_output = generate_output(reranked_results, generated_queries)

    print(final_output)

示例执行结果如下，

生成的查询：
1. "气候变化对生态系统的影响"
2. "气候变化对水资源的影响"
3. "气候变化对气候系统的影响"
4. "气候变化对人类健康的影响"
初始各个查询搜索结果排名：
查询： '1. "气候变化对生态系统的影响"': {'doc6': 0.84, 'doc9': 0.81, 'doc5': 0.77, 'doc2': 0.75, 'doc8': 0.74}
查询： '2. "气候变化对水资源的影响"': {'doc8': 0.83, 'doc9': 0.82, 'doc7': 0.82, 'doc2': 0.77, 'doc4': 0.72}
查询： '3. "气候变化对气候系统的影响"': {'doc5': 0.87, 'doc8': 0.87, 'doc1': 0.84, 'doc4': 0.74, 'doc3': 0.73}
查询： '4. "气候变化对人类健康的影响"': {'doc8': 0.89, 'doc10': 0.88, 'doc4': 0.84, 'doc5': 0.76}
根据查询 '1. "气候变化对生态系统的影响"' 中的排名更新 doc6 的得分，从 0 更新为 0.016666666666666666 
根据查询 '1. "气候变化对生态系统的影响"' 中的排名更新 doc9 的得分，从 0 更新为 0.01639344262295082 
根据查询 '1. "气候变化对生态系统的影响"' 中的排名更新 doc5 的得分，从 0 更新为 0.016129032258064516 
根据查询 '1. "气候变化对生态系统的影响"' 中的排名更新 doc2 的得分，从 0 更新为 0.015873015873015872 
根据查询 '1. "气候变化对生态系统的影响"' 中的排名更新 doc8 的得分，从 0 更新为 0.015625 
根据查询 '2. "气候变化对水资源的影响"' 中的排名更新 doc8 的得分，从 0.015625 更新为 0.03229166666666666 
根据查询 '2. "气候变化对水资源的影响"' 中的排名更新 doc9 的得分，从 0.01639344262295082 更新为 0.03278688524590164 
根据查询 '2. "气候变化对水资源的影响"' 中的排名更新 doc7 的得分，从 0 更新为 0.016129032258064516 
根据查询 '2. "气候变化对水资源的影响"' 中的排名更新 doc2 的得分，从 0.015873015873015872 更新为 0.031746031746031744 
根据查询 '2. "气候变化对水资源的影响"' 中的排名更新 doc4 的得分，从 0 更新为 0.015625 
根据查询 '3. "气候变化对气候系统的影响"' 中的排名更新 doc5 的得分，从 0.016129032258064516 更新为 0.03279569892473118 
根据查询 '3. "气候变化对气候系统的影响"' 中的排名更新 doc8 的得分，从 0.03229166666666666 更新为 0.04868510928961749 
根据查询 '3. "气候变化对气候系统的影响"' 中的排名更新 doc1 的得分，从 0 更新为 0.016129032258064516 
根据查询 '3. "气候变化对气候系统的影响"' 中的排名更新 doc4 的得分，从 0.015625 更新为 0.03149801587301587 
根据查询 '3. "气候变化对气候系统的影响"' 中的排名更新 doc3 的得分，从 0 更新为 0.015625 
根据查询 '4. "气候变化对人类健康的影响"' 中的排名更新 doc8 的得分，从 0.04868510928961749 更新为 0.06535177595628415 
根据查询 '4. "气候变化对人类健康的影响"' 中的排名更新 doc10 的得分，从 0 更新为 0.01639344262295082 
根据查询 '4. "气候变化对人类健康的影响"' 中的排名更新 doc4 的得分，从 0.03149801587301587 更新为 0.04762704813108039 
根据查询 '4. "气候变化对人类健康的影响"' 中的排名更新 doc5 的得分，从 0.03279569892473118 更新为 0.04866871479774705 
最终重新排序的结果： {'doc8': 0.06535177595628415, 'doc5': 0.04866871479774705, 'doc4': 0.04762704813108039, 'doc9': 0.03278688524590164, 'doc2': 0.031746031746031744, 'doc6': 0.016666666666666666, 'doc10': 0.01639344262295082, 'doc7': 0.016129032258064516, 'doc1': 0.016129032258064516, 'doc3': 0.015625}
根据查询 ['1. "气候变化对生态系统的影响"', '2. "气候变化对水资源的影响"', '3. "气候变化对气候系统的影响"', '4. "气候变化对人类健康的影响"'] 和重新排序的文档的最终输出：['doc8', 'doc5', 'doc4', 'doc9', 'doc2', 'doc6', 'doc10', 'doc7', 'doc1', 'doc3']

完成！