18种RAG技术大比拼：谁才是检索增强生成的最佳选择？ 原创 精华

# 定义 PDF 文件路径
pdf_path = "data/AI_information.pdf"

# 提取文本并分割成文本块
extracted_text = extract_text_from_pdf(pdf_path)
text_chunks = chunk_text(extracted_text, chunk_size=1000, overlap=200)

# 打印文本块数量
print("Number of text chunks:", len(text_chunks))
1.
2.
3.
4.
5.
6.
7.
8.
9.

# 将提取的文本分割成句子
sentences = extracted_text.split(". ")

# 为每个句子生成嵌入向量
embeddings = [get_embedding(sentence) for sentence in sentences]

# 计算连续句子之间的相似度
similarities = [cosine_similarity(embeddings[i], embeddings[i + 1]) for i in range(len(embeddings) - 1)]

# 使用百分位法计算断点
breakpoints = compute_breakpoints(similarities, method="percentile", threshold=90)

# 根据断点将句子分组成语义块
text_chunks = split_into_chunks(sentences, breakpoints)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.

# 定义上下文增强检索函数
def context_enriched_search(query, text_chunks, embeddings, k=1, context_size=1):
    query_embedding = create_embeddings(query).data[0].embedding
    similarity_scores = []

    for i, chunk_embedding in enumerate(embeddings):
        similarity_score = cosine_similarity(np.array(query_embedding), np.array(chunk_embedding.embedding))
        similarity_scores.append((i, similarity_score))

    similarity_scores.sort(key=lambda x: x[1], reverse=True)
    top_index = similarity_scores[0][0]

    start = max(0, top_index - context_size)
    end = min(len(text_chunks), top_index + context_size + 1)

    return [text_chunks[i] for i in range(start, end)]
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.

# 为每个文本块生成描述性标题
text_chunks_with_headers = chunk_text_with_headers(extracted_text, chunk_size=1000, overlap=200)

# 打印一个样本以查看效果
print("Sample Chunk with Header:")
print("Header:", text_chunks_with_headers[0]['header'])
print("Content:", text_chunks_with_headers[0]['text'])
1.
2.
3.
4.
5.
6.
7.

# 处理文档，提取文本、创建文本块、生成问题并构建向量库
text_chunks, vector_store = process_document(
    pdf_path,
    chunk_size=1000,
    chunk_overlap=200,
    questions_per_chunk=3
)
1.
2.
3.
4.
5.
6.
7.

# 查询重写
rewritten_query = rewrite_query(query)

# 后退提示
step_back_query = generate_step_back_query(query)

# 子查询分解
sub_queries = decompose_query(query, num_subqueries=4)
1.
2.
3.
4.
5.
6.
7.
8.

# 定义重排序函数
def rag_with_reranking(query, vector_store, reranking_method="llm", top_n=3, model="meta-llama/Llama-3.2-3B-Instruct"):
    query_embedding = create_embeddings(query)
    initial_results = vector_store.similarity_search(query_embedding, k=10)

    if reranking_method == "llm":
        reranked_results = rerank_with_llm(query, initial_results, top_n=top_n)
    else:
        reranked_results = initial_results[:top_n]

    context = "\n\n===\n\n".join([result["text"] for result in reranked_results])
    response = generate_response(query, context, model)

    return {
        "query": query,
        "reranking_method": reranking_method,
        "initial_results": initial_results[:top_n],
        "reranked_results": reranked_results,
        "context": context,
        "response": response
    }
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.

# 运行 RAG 并应用相关段落提取
rse_result = rag_with_rse(pdf_path, query)
1.
2.

# 定义上下文压缩函数
def rag_with_compression(pdf_path, query, k=10, compression_type="selective", model="meta-llama/Llama-3.2-3B-Instruct"):
    vector_store = process_document(pdf_path)
    results = vector_store.similarity_search(create_embeddings(query), k=k)
    retrieved_chunks = [r["text"] for r in results]

    compressed = batch_compress_chunks(retrieved_chunks, query, compression_type, model)
    compressed_chunks, compression_ratios = zip([(c, r) for c, r in compressed if c.strip()] or [(chunk, 0.0) for chunk in retrieved_chunks])

    context = "\n\n---\n\n".join(compressed_chunks)
    response = generate_response(query, context, model)

    return {
        "query": query,
        "original_chunks": retrieved_chunks,
        "compressed_chunks": compressed_chunks,
        "compression_ratios": compression_ratios,
        "context_length_reduction": f"{sum(compression_ratios)/len(compression_ratios):.2f}%",
        "response": response
    }
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.

# 定义完整的 RAG 流程，包含反馈机制
def full_rag_workflow(pdf_path, query, feedback_data=None, feedback_file="feedback_data.json", fine_tune=False):
    if feedback_data isNone:
        feedback_data = load_feedback_data(feedback_file)
        print(f"Loaded {len(feedback_data)} feedback entries from {feedback_file}")

    chunks, vector_store = process_document(pdf_path)

    if fine_tune and feedback_data:
        vector_store = fine_tune_index(vector_store, chunks, feedback_data)

    result = rag_with_feedback_loop(query, vector_store, feedback_data)

    print("\n=== Would you like to provide feedback on this response? ===")
    print("Rate relevance (1-5, with 5 being most relevant):")
    relevance = input()

    print("Rate quality (1-5, with 5 being highest quality):")
    quality = input()

    print("Any comments? (optional, press Enter to skip)")
    comments = input()

    feedback = get_user_feedback(
        query=query,
        respnotallow=result["response"],
        relevance=int(relevance),
        quality=int(quality),
        comments=comments
    )

    store_feedback(feedback, feedback_file)
    print("Feedback recorded. Thank you!")

    return result
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.

# 定义自适应 RAG 函数
def rag_with_adaptive_retrieval(pdf_path, query, k=4, user_cnotallow=None):
    chunks, vector_store = process_document(pdf_path)
    query_type = classify_query(query)
    print(f"Query classified as: {query_type}")

    retrieved_docs = adaptive_retrieval(query, vector_store, k, user_context)
    response = generate_response(query, retrieved_docs, query_type)

    return {
        "query": query,
        "query_type": query_type,
        "retrieved_documents": retrieved_docs,
        "response": response
    }
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.

# 运行自我 RAG
result = self_rag(query, vector_store)
1.
2.

# 构建知识图谱
def build_knowledge_graph(chunks):
    graph, texts = nx.Graph(), [c["text"] for c in chunks]
    embeddings = create_embeddings(texts)

    for i, (chunk, emb) in enumerate(zip(chunks, embeddings)):
        graph.add_node(i, text=chunk["text"], cnotallow=extract_concepts(chunk["text"]), embedding=emb)

    for i, j in ((i, j) for i in range(len(chunks)) for j in range(i + 1, len(chunks))):
        if shared_concepts := set(graph.nodes[i]["concepts"]) & set(graph.nodes[j]["concepts"]):
            sim = np.dot(embeddings[i], embeddings[j]) / (np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[j]))
            weight = 0.7 * sim + 0.3 * (len(shared_concepts) / min(len(graph.nodes[i]["concepts"]), len(graph.nodes[j]["concepts"])))
            if weight > 0.6:
                graph.add_edge(i, j, weight=weight, similarity=sim, shared_cnotallow=list(shared_concepts))

    print(f"Graph built: {graph.number_of_nodes()} nodes, {graph.number_of_edges()} edges")
    return graph, embeddings
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.

# 定义层次化 RAG 函数
def hierarchical_rag(query, pdf_path, chunk_size=1000, chunk_overlap=200, k_summaries=3, k_chunks=5, regenerate=False):
    summary_store_file = f"{os.path.basename(pdf_path)}_summary_store.pkl"
    detailed_store_file = f"{os.path.basename(pdf_path)}_detailed_store.pkl"

    if regenerate ornot os.path.exists(summary_store_file) ornot os.path.exists(detailed_store_file):
        print("Processing document and creating vector stores...")
        summary_store, detailed_store = process_document_hierarchically(pdf_path, chunk_size, chunk_overlap)

        with open(summary_store_file, 'wb') as f:
            pickle.dump(summary_store, f)
        with open(detailed_store_file, 'wb') as f:
            pickle.dump(detailed_store, f)
    else:
        print("Loading existing vector stores...")
        with open(summary_store_file, 'rb') as f:
            summary_store = pickle.load(f)
        with open(detailed_store_file, 'rb') as f:
            detailed_store = pickle.load(f)

    retrieved_chunks = retrieve_hierarchically(query, summary_store, detailed_store, k_summaries, k_chunks)
    response = generate_response(query, retrieved_chunks)

    return {
        "query": query,
        "response": response,
        "retrieved_chunks": retrieved_chunks,
        "summary_count": len(summary_store.texts),
        "detailed_count": len(detailed_store.texts)
    }
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.

# 定义 HyDE RAG 函数
def hyde_rag(query, vector_store, k=5, should_generate_respnotallow=True):
    print(f"\n=== Processing query with HyDE: {query} ===\n")

    print("Generating hypothetical document...")
    hypothetical_doc = generate_hypothetical_document(query)
    print(f"Generated hypothetical document of {len(hypothetical_doc)} characters")

    print("Creating embedding for hypothetical document...")
    hypothetical_embedding = create_embeddings([hypothetical_doc])[0]

    print(f"Retrieving {k} most similar chunks...")
    retrieved_chunks = vector_store.similarity_search(hypothetical_embedding, k=k)

    results = {
        "query": query,
        "hypothetical_document": hypothetical_doc,
        "retrieved_chunks": retrieved_chunks
    }

    if should_generate_response:
        print("Generating final response...")
        response = generate_response(query, retrieved_chunks)
        results["response"] = response

    return results
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.

# 定义融合检索函数
def fusion_retrieval(query, chunks, vector_store, bm25_index, k=5, alpha=0.5):
    query_embedding = create_embeddings(query)
    vector_results = {r["metadata"]["index"]: r["similarity"] for r in vector_store.similarity_search_with_scores(query_embedding, len(chunks))}
    bm25_results = {r["metadata"]["index"]: r["bm25_score"] for r in bm25_search(bm25_index, chunks, query, len(chunks))}

    all_docs = vector_store.get_all_documents()
    scores = [(i, alpha * vector_results.get(i, 0) + (1 - alpha) * bm25_results.get(i, 0)) for i in range(len(all_docs))]
    top_docs = sorted(scores, key=lambda x: x[1], reverse=True)[:k]

    return [{"text": all_docs[i]["text"], "metadata": all_docs[i]["metadata"], "score": s} for i, s in top_docs]
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.

# 处理多模态文档
def process_document(pdf_path, chunk_size=1000, chunk_overlap=200):
    image_dir = "extracted_images"
    os.makedirs(image_dir, exist_ok=True)

    text_data, image_paths = extract_content_from_pdf(pdf_path, image_dir)
    chunked_text = chunk_text(text_data, chunk_size, chunk_overlap)
    image_data = process_images(image_paths)

    all_items = chunked_text + image_data
    contents = [item["content"] for item in all_items]
    embeddings = create_embeddings(contents)

    vector_store = MultiModalVectorStore()
    vector_store.add_items(all_items, embeddings)

    doc_info = {
        "text_count": len(chunked_text),
        "image_count": len(image_data),
        "total_items": len(all_items),
    }

    print(f"Added {len(all_items)} items to vector store ({len(chunked_text)} text chunks, {len(image_data)} image captions)")
    return vector_store, doc_info
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.

# 运行纠正性 RAG
crag_result = rag_with_compression(pdf_path, query, compression_type="selective")
1.
2.