2W8000字深度剖析25种RAG变体：全网最全没有之一 精华

from transformers import AutoTokenizer, AutoModel
import torch

tokenizer = AutoTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
model = AutoModel.from_pretrained("facebook/dpr-question_encoder-single-nq-base")

def encode_query(query: str) -> torch.Tensor:
    inputs = tokenizer(query, return_tensors="pt", truncatinotallow=True, max_length=128)
    return model(**inputs).pooler_output

import nltk
nltk.download('punkt')
from nltk.tokenize import sent_tokenize

def segment_document(document: str, window_size: int = 5, overlap: int = 2) -> list:
    sentences = sent_tokenize(document)
    return [" ".join(sentences[i:i+window_size]) for i in range(0, len(sentences), window_size - overlap)]

import faiss, numpy as np

def build_faiss_index(embeddings: torch.Tensor) -> faiss.IndexFlatL2:
    emb_np = embeddings.detach().cpu().numpy().astype('float32')
    d = emb_np.shape[1]
    index = faiss.IndexFlatL2(d)
    index.add(emb_np)
    return index

def construct_prompt(query: str, retrieved_chunks: list) -> str:
    delimiter = "[DOC]"
    context = " ".join([f"{delimiter} {chunk} {delimiter}" for chunk in retrieved_chunks])
    return f"{query}\n{context}"

from transformers import AutoModelForCausalLM

gen_tokenizer = AutoTokenizer.from_pretrained("gpt2")
gen_model = AutoModelForCausalLM.from_pretrained("gpt2")

def generate_response(prompt: str, max_length: int = 200) -> str:
    inputs = gen_tokenizer.encode(prompt, return_tensors="pt", truncatinotallow=True, max_length=1024)
    outputs = gen_model.generate(inputs, max_length=max_length, num_beams=5, early_stopping=True)
    return gen_tokenizer.decode(outputs[0], skip_special_tokens=True)

def evaluate_confidence(output_ids: torch.Tensor) -> float:
    outputs = gen_model(output_ids, labels=output_ids)
    return torch.exp(-outputs.loss).item()

def iterative_crag(prompt: str, threshold: float = 0.8, max_iters: int = 3) -> str:
    current_prompt = prompt
    for _ in range(max_iters):
        response = generate_response(current_prompt)
        output_ids = gen_tokenizer.encode(response, return_tensors="pt")
        conf = evaluate_confidence(output_ids)
        if conf >= threshold:
            return response
        current_prompt += "\n[REFINE]"
    return response

def generate_candidates(prompt: str, num_candidates: int = 3) -> list:
    inputs = gen_tokenizer.encode(prompt, return_tensors="pt", truncatinotallow=True, max_length=1024)
    candidates_ids = gen_model.generate(inputs, max_length=200, num_return_sequences=num_candidates, num_beams=5)
    return [gen_tokenizer.decode(c, skip_special_tokens=True) for c in candidates_ids]

def select_best_candidate(candidates: list) -> str:
    best_candidate = None
    best_score = -float('inf')
    for candidate in candidates:
        inputs = gen_tokenizer.encode(candidate, return_tensors="pt", truncatinotallow=True, max_length=1024)
        outputs = gen_model(inputs, labels=inputs)
        score = torch.exp(-outputs.loss).item()
        if score > best_score:
            best_score = score
            best_candidate = candidate
    return best_candidate

def speculative_rag(prompt: str) -> str:
    candidates = generate_candidates(prompt, num_candidates=3)
    return select_best_candidate(candidates)

# Placeholder for BM25 retrieval (e.g., using rank_bm25 library)
def bm25_retrieve(query: str, corpus: list, k: int = 5):
    from rank_bm25 import BM25Okapi
    tokenized_corpus = [doc.split() for doc in corpus]
    bm25 = BM25Okapi(tokenized_corpus)
    tokenized_query = query.split()
    scores = bm25.get_scores(tokenized_query)
    top_k = np.argsort(scores)[-k:]
    return top_k

def hybrid_retrieve(query: str, dense_index: faiss.IndexFlatL2, corpus: list, k: int = 5):
    query_emb = encode_query(query)
    _, indices_dense = retrieve_documents(query_emb, dense_index, k)
    indices_bm25 = bm25_retrieve(query, corpus, k)
    # Merge indices (simple union, can use weighted merging)
    merged_indices = list(set(indices_dense[0]).union(set(indices_bm25)))
    return merged_indices

import requests

def query_reformulator(raw_query: str) -> str:
    response = requests.post("http://agent-service/reformulate", jsnotallow={"query": raw_query})
    return response.json().get("reformulated_query", raw_query)

def document_retriever(query: str) -> list:
    response = requests.post("http://agent-service/retrieve", jsnotallow={"query": query})
    return response.json().get("documents", [])

def response_synthesizer(original_query: str, documents: list) -> str:
    payload = {"query": original_query, "documents": documents}
    response = requests.post("http://agent-service/synthesize", jsnotallow=payload)
    return response.json().get("response", "")

def agentic_rag(raw_query: str) -> str:
    reformulated_query = query_reformulator(raw_query)
    docs = document_retriever(reformulated_query)
    return response_synthesizer(raw_query, docs)

def self_rag(prompt: str, iterations: int = 3) -> str:
    current_prompt = prompt
    for _ in range(iterations):
        response = generate_response(current_prompt, gen_model, gen_tokenizer)
        response_emb = encode_query(response)
        additional_context = " [ADDITIONAL EVIDENCE REQUIRED]"
        current_prompt = f"{prompt}\n{additional_context}"
    return response

def gating_decision(query: str, threshold: float = 0.85) -> bool:
    query_emb = encode_query(query)
    confidence = query_emb.norm().item() / 100
    return confidence < threshold

def adaptive_rag(query: str, dense_index: faiss.IndexFlatL2) -> str:
    if gating_decision(query):
        distances, indices = retrieve_documents(encode_query(query), dense_index, k=5)
        retrieved_docs = [chunks[i] for i in indices[0]]
        prompt = construct_prompt(query, retrieved_docs)
    else:
        prompt = query
    return generate_response(prompt, gen_model, gen_tokenizer)

def refeed_rag(query: str, dense_index: faiss.IndexFlatL2) -> str:
    initial_response = generate_response(query, gen_model, gen_tokenizer)
    feedback_query = f"{query} {initial_response}"
    distances, indices = retrieve_documents(encode_query(feedback_query), dense_index, k=5)
    additional_context = " ".join([chunks[i] for i in indices[0]])
    new_prompt = f"{query}\n[REFEED CONTEXT] {additional_context}"
    return generate_response(new_prompt, gen_model, gen_tokenizer)

def cluster_chunks(embeddings: torch.Tensor, num_clusters: int = 10):
    from sklearn.cluster import KMeans
    emb_np = embeddings.detach().cpu().numpy()
    kmeans = KMeans(n_clusters=num_clusters)
    labels = kmeans.fit_predict(emb_np)
    return labels, kmeans.cluster_centers_


def build_tree(chunks: list, labels: list):
    tree = {}
    for label in set(labels):
        tree[label] = [chunks[i] for i, lab in enumerate(labels) if lab == label]
    return tree


# 在查询时：
def tree_retrieve(query: str, tree: dict):
    aspects = query.split()  # 用于适当分解的占位符
    retrieved = []
    for aspect in aspects:
        for key, docs in tree.items():
            if aspect.lower() in" ".join(docs).lower():
                retrieved.extend(docs)
    return list(set(retrieved))


# 使用方法：
labels, _ = cluster_chunks(chunk_embeddings, num_clusters=10)
raptor_tree = build_tree(chunks, labels)
tree_docs = tree_retrieve(query, raptor_tree)
raptor_prompt = construct_prompt(query, tree_docs)
raptor_response = generate_response(raptor_prompt, gen_model, gen_tokenizer)

def generate_with_actions(prompt: str) -> str:
    inputs = gen_tokenizer.encode(prompt, return_tensors="pt", truncatinotallow=True, max_length=300)
    outputs = gen_model.generate(inputs, max_length=300, num_beams=5)
    return gen_tokenizer.decode(outputs[0], skip_special_tokens=True)

def replug_rag(query: str, dense_index: faiss.IndexFlatL2) -> str:
    query_emb = encode_query(query)
    _, indices = retrieve_documents(query_emb, dense_index, k=5)
    retrieved = [chunks[i] for i in indices[0]]
    enriched_prompt = construct_prompt(query, retrieved)
    return generate_response(enriched_prompt, gen_model, gen_tokenizer)

# 假设memory_module是一个预训练模块，具有generate_clue方法。
def memo_rag(query: str, dense_index: faiss.IndexFlatL2, memory_module) -> str:
    clue = memory_module.generate_clue(query)
    feedback_query = f"{query} {clue}"
    _, indices = retrieve_documents(encode_query(feedback_query), dense_index, k=5)
    retrieved = [chunks[i] for i in indices[0]]
    prompt = construct_prompt(query, retrieved)
    return generate_response(prompt, gen_model, gen_tokenizer)

def scaled_dot_product_attention(Q, K, V, mask=None):
    import math
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    attention_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, V)
    return output, attention_weights

def retro_retrieve(block_embedding, index, k=5):
    distances, indices = index.search(block_embedding.detach().cpu().numpy().astype('float32'), k)
    return indices

output = scaled_dot_product_attention(Q, K, V)[0]

def auto_rag(prompt: str, max_iterations: int = 3, threshold: float = 0.8) -> str:
    current_prompt = prompt
    for i in range(max_iterations):
        response = generate_response(current_prompt, gen_model, gen_tokenizer)
        output_ids = gen_tokenizer.encode(response, return_tensors="pt")
        conf = evaluate_confidence(output_ids)
        if conf >= threshold:
            break
        current_prompt += "\n[RETRIEVE_MORE]"
    return response

def mcts_retrieval(query: str, index: faiss.IndexFlatL2, iterations: int = 10):
    best_set = None
    best_score = -float('inf')
    for _ in range(iterations):
        candidate_set = sample_candidate_set(index)
        quality = evaluate_candidate_quality(candidate_set, query)
        cost = evaluate_candidate_cost(candidate_set)
        score = quality - cost
        if score > best_score:
            best_score = score
            best_set = candidate_set
    return best_set

def edge_retrieve(query: str, local_index: faiss.IndexFlatL2) -> list:
    q_emb = encode_query(query)
    _, indices = retrieve_documents(q_emb, local_index, k=5)
    return [chunks[i] for i in indices[0]]

def collaborative_retrieve(query: str, edge_indices: list) -> list:
    results = []
    for idx in edge_indices:
        results.extend(edge_retrieve(query, idx))
    return rank_results(results)

def apply_rules(query: str, documents: list, keyword: str = "RAG") -> list:
    return [doc for doc in documents if keyword.lower() in doc.lower()]

def rule_rag(query: str, dense_index: faiss.IndexFlatL2) -> str:
    _, indices = retrieve_documents(encode_query(query), dense_index, k=10)
    raw_docs = [chunks[i] for i in indices[0]]
    filtered_docs = apply_rules(query, raw_docs, keyword="RAG")
    prompt = construct_prompt(query, filtered_docs)
    return generate_response(prompt, gen_model, gen_tokenizer)

def conversational_rag(dialogue_history: list, query: str, dense_index: faiss.IndexFlatL2) -> str:
    context = " ".join(dialogue_history)
    full_query = f"{context} {query}"
    _, indices = retrieve_documents(encode_query(full_query), dense_index, k=5)
    retrieved = [chunks[i] for i in indices[0]]
    prompt = construct_prompt(full_query, retrieved)
    return generate_response(prompt, gen_model, gen_tokenizer)

def iterative_rag(query: str, dense_index: faiss.IndexFlatL2, max_iters: int = 3) -> str:
    current_prompt = query
    for _ in range(max_iters):
        response = generate_response(current_prompt, gen_model, gen_tokenizer)
        if "unclear" not in response:
            return response
        _, indices = retrieve_documents(encode_query(response), dense_index, k=5)
        additional_context = " ".join([chunks[i] for i in indices[0]])
        current_prompt = f"{query}\n{additional_context}"
    return response

def decompose_query(query: str) -> list:
    # 为演示目的，按标点分割；如有需要，可替换为高级主题建模方法
    return query.split(";")

def tree_retrieve(aspects: list, dense_index: faiss.IndexFlatL2) -> list:
    tree_results = []
    for aspect in aspects:
        _, indices = retrieve_documents(encode_query(aspect), dense_index, k=3)
        tree_results.extend([chunks[i] for i in indices[0]])
    return list(set(tree_results))

def tree_rag(query: str, dense_index: faiss.IndexFlatL2) -> str:
    aspects = decompose_query(query)
    tree_docs = tree_retrieve(aspects, dense_index)
    prompt = construct_prompt(query, tree_docs)
    return generate_response(prompt, gen_model, gen_tokenizer)

def detect_unknown_terms(text: str) -> list:
    return [word for word in text.split() if word.isupper()]

def construct_transkg(terms: list) -> dict:
    return {term: f"translated_{term}" for term in terms}

def crat_translate(text: str) -> str:
    unknown_terms = detect_unknown_terms(text)
    transkg = construct_transkg(unknown_terms)
    for term, translation in transkg.items():
        text = text.replace(term, translation)
    return text

import spacy, networkx as nx
nlp = spacy.load("en_core_web_sm")

def extract_entities_relations(text: str) -> (list, list):
    doc = nlp(text)
    entities = [ent.text for ent in doc.ents]
    relations = [(entities[i], entities[j]) for i in range(len(entities)) for j in range(i+1, len(entities))]
    return entities, relations

def build_graph(documents: list) -> nx.Graph:
    G = nx.Graph()
    for doc in documents:
        entities, relations = extract_entities_relations(doc)
        for entity in entities:
            G.add_node(entity)
        for (e1, e2) in relations:
            G.add_edge(e1, e2)
    return G

def graph_based_prompt(query: str, G: nx.Graph) -> str:
    relevant_nodes = [node for node in G.nodes if node.lower() in query.lower()]
    graph_context = " ".join(relevant_nodes)
    return f"{query}\n[GRAPH_CONTEXT] {graph_context}"