RAG 论文综述：Retrieval-Augmented Generation 的演进 — 从 RETRO 到 Self-RAG、Corrective-RAG

import torch
import numpy as np
from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer
from transformers import DPRContextEncoder, DPRContextEncoderTokenizer

class DPRRetriever:
    """基于 DPR 的 Dense Passage Retrieval 实现"""

    def __init__(self, model_name="facebook/dpr-question_encoder-single-nq-base"):
        self.q_encoder = DPRQuestionEncoder.from_pretrained(model_name)
        self.q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(model_name)

        ctx_model = "facebook/dpr-ctx_encoder-single-nq-base"
        self.ctx_encoder = DPRContextEncoder.from_pretrained(ctx_model)
        self.ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained(ctx_model)

        self.document_embeddings = None
        self.documents = []

    def encode_documents(self, documents: list[str]) -> np.ndarray:
        """将文档语料转换为向量表示"""
        self.documents = documents
        embeddings = []

        for doc in documents:
            inputs = self.ctx_tokenizer(
                doc, return_tensors="pt",
                max_length=256, truncation=True, padding=True
            )
            with torch.no_grad():
                output = self.ctx_encoder(**inputs)
            embeddings.append(output.pooler_output.numpy())

        self.document_embeddings = np.vstack(embeddings)
        # 进行 L2 归一化
        norms = np.linalg.norm(self.document_embeddings, axis=1, keepdims=True)
        self.document_embeddings = self.document_embeddings / norms
        return self.document_embeddings

    def retrieve(self, query: str, top_k: int = 5) -> list[dict]:
        """检索与查询最相关的前 k 个文档"""
        inputs = self.q_tokenizer(
            query, return_tensors="pt",
            max_length=64, truncation=True, padding=True
        )
        with torch.no_grad():
            q_embedding = self.q_encoder(**inputs).pooler_output.numpy()

        q_embedding = q_embedding / np.linalg.norm(q_embedding)
        scores = np.dot(self.document_embeddings, q_embedding.T).squeeze()
        top_indices = np.argsort(scores)[::-1][:top_k]

        results = []
        for idx in top_indices:
            results.append({
                "document": self.documents[idx],
                "score": float(scores[idx]),
                "index": int(idx)
            })
        return results


# 使用示例
retriever = DPRRetriever()
corpus = [
    "RAG는 검색과 생성을 결합한 모델이다.",
    "Transformer는 Self-Attention 메커니즘을 사용한다.",
    "BERT는 양방향 사전학습 언어 모델이다.",
    "DPR은 밀집 벡터를 사용하여 패시지를 검색한다.",
]
retriever.encode_documents(corpus)
results = retriever.retrieve("RAG에서 문서 검색은 어떻게 동작하나요?")
for r in results:
    print(f"[Score: {r['score']:.4f}] {r['document']}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ChunkedCrossAttention(nn.Module):
    """RETRO 风格的 Chunked Cross-Attention 实现"""

    def __init__(self, d_model: int = 512, n_heads: int = 8, chunk_size: int = 64):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.chunk_size = chunk_size

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(
        self,
        hidden_states: torch.Tensor,
        retrieved_chunks: torch.Tensor
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: (B, seq_len, d_model) - 解码器隐藏状态
            retrieved_chunks: (B, n_chunks, chunk_len, d_model) - 检索到的相邻文本块
        """
        B, seq_len, D = hidden_states.shape
        n_chunks = seq_len // self.chunk_size

        # 按 chunk 切分序列
        h_chunks = hidden_states[:, :n_chunks * self.chunk_size].reshape(
            B, n_chunks, self.chunk_size, D
        )

        # 对每个 chunk 与检索到的相邻块执行 Cross-Attention
        Q = self.W_q(h_chunks)  # (B, n_chunks, chunk_size, D)
        K = self.W_k(retrieved_chunks)  # (B, n_chunks, chunk_len, D)
        V = self.W_v(retrieved_chunks)

        # 拆分为多头
        Q = Q.reshape(B, n_chunks, self.chunk_size, self.n_heads, self.d_k).permute(0, 1, 3, 2, 4)
        K = K.reshape(B, n_chunks, -1, self.n_heads, self.d_k).permute(0, 1, 3, 2, 4)
        V = V.reshape(B, n_chunks, -1, self.n_heads, self.d_k).permute(0, 1, 3, 2, 4)

        # Scaled Dot-Product Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        attn_weights = F.softmax(scores, dim=-1)
        attn_output = torch.matmul(attn_weights, V)

        # 合并多头并做输出投影
        attn_output = attn_output.permute(0, 1, 3, 2, 4).reshape(
            B, n_chunks, self.chunk_size, D
        )
        attn_output = self.W_o(attn_output)

        # 残差连接与层归一化
        output = self.layer_norm(h_chunks + attn_output)
        output = output.reshape(B, n_chunks * self.chunk_size, D)

        # 恢复剩余 token（当序列长度无法被 chunk_size 整除时）
        if seq_len > n_chunks * self.chunk_size:
            remainder = hidden_states[:, n_chunks * self.chunk_size:]
            output = torch.cat([output, remainder], dim=1)

        return output


# RETRO 风格检索流水线示例
cca = ChunkedCrossAttention(d_model=512, n_heads=8, chunk_size=64)
hidden = torch.randn(2, 256, 512)  # 배치 2, 시퀀스 256
retrieved = torch.randn(2, 4, 32, 512)  # 4 청크, 각 32 토큰
output = cca(hidden, retrieved)
print(f"Input shape: {hidden.shape} -> Output shape: {output.shape}")

from dataclasses import dataclass
from enum import Enum
from typing import Optional

class RetrieveDecision(Enum):
    YES = "yes"
    NO = "no"
    CONTINUE = "continue"

class RelevanceScore(Enum):
    RELEVANT = "relevant"
    IRRELEVANT = "irrelevant"

class SupportScore(Enum):
    FULLY_SUPPORTED = "fully_supported"
    PARTIALLY_SUPPORTED = "partially_supported"
    NO_SUPPORT = "no_support"

@dataclass
class ReflectionResult:
    retrieve: RetrieveDecision
    relevance: Optional[RelevanceScore] = None
    support: Optional[SupportScore] = None
    utility: Optional[int] = None  # 1-5 分

class SelfRAGPipeline:
    """Self-RAG 风格的自适应检索-生成流水线"""

    def __init__(self, generator, retriever, reflection_model):
        self.generator = generator
        self.retriever = retriever
        self.reflection_model = reflection_model

    def should_retrieve(self, query: str, partial_output: str = "") -> RetrieveDecision:
        """自主判断是否需要检索的反思步骤"""
        prompt = (
            f"Query: {query}\n"
            f"Partial output: {partial_output}\n"
            "Does this query require external knowledge retrieval? "
            "Answer: yes, no, or continue"
        )
        decision = self.reflection_model.predict(prompt)
        return RetrieveDecision(decision.strip().lower())

    def evaluate_relevance(self, query: str, document: str) -> RelevanceScore:
        """评估检索文档的相关性（模拟 ISREL token）"""
        prompt = (
            f"Query: {query}\n"
            f"Document: {document}\n"
            "Is this document relevant to answering the query? "
            "Answer: relevant or irrelevant"
        )
        score = self.reflection_model.predict(prompt)
        return RelevanceScore(score.strip().lower())

    def evaluate_support(
        self, query: str, document: str, response: str
    ) -> SupportScore:
        """评估生成结果的依据充分性（模拟 ISSUP token）"""
        prompt = (
            f"Query: {query}\n"
            f"Document: {document}\n"
            f"Response: {response}\n"
            "Is the response supported by the document? "
            "Answer: fully_supported, partially_supported, or no_support"
        )
        score = self.reflection_model.predict(prompt)
        return SupportScore(score.strip().lower())

    def generate_with_reflection(self, query: str) -> dict:
        """执行完整的 Self-RAG 流水线"""
        # 第 1 步：判断是否需要检索
        retrieve_decision = self.should_retrieve(query)

        if retrieve_decision == RetrieveDecision.NO:
            # 无需检索 —— 直接生成
            response = self.generator.generate(query)
            return {
                "response": response,
                "retrieved": False,
                "reflection": ReflectionResult(retrieve=RetrieveDecision.NO)
            }

        # 第 2 步：检索文档
        documents = self.retriever.retrieve(query, top_k=5)

        # 第 3 步：通过相关性评估过滤文档
        relevant_docs = []
        for doc in documents:
            relevance = self.evaluate_relevance(query, doc["text"])
            if relevance == RelevanceScore.RELEVANT:
                relevant_docs.append(doc)

        if not relevant_docs:
            # 没有相关文档 —— 不检索直接生成
            response = self.generator.generate(query)
            return {
                "response": response,
                "retrieved": True,
                "relevant_docs": 0,
                "reflection": ReflectionResult(
                    retrieve=RetrieveDecision.YES,
                    relevance=RelevanceScore.IRRELEVANT
                )
            }

        # 第 4 步：生成候选回答并评估
        best_response = None
        best_score = -1

        for doc in relevant_docs:
            context = f"Context: {doc['text']}\nQuery: {query}"
            candidate = self.generator.generate(context)

            support = self.evaluate_support(query, doc["text"], candidate)
            # 计算支持度得分
            support_score = {
                SupportScore.FULLY_SUPPORTED: 3,
                SupportScore.PARTIALLY_SUPPORTED: 1,
                SupportScore.NO_SUPPORT: 0
            }.get(support, 0)

            if support_score > best_score:
                best_score = support_score
                best_response = candidate
                best_support = support

        return {
            "response": best_response,
            "retrieved": True,
            "relevant_docs": len(relevant_docs),
            "reflection": ReflectionResult(
                retrieve=RetrieveDecision.YES,
                relevance=RelevanceScore.RELEVANT,
                support=best_support,
                utility=min(best_score + 2, 5)
            )
        }

from dataclasses import dataclass
from enum import Enum
import numpy as np

class ConfidenceLevel(Enum):
    CORRECT = "correct"
    INCORRECT = "incorrect"
    AMBIGUOUS = "ambiguous"

@dataclass
class EvaluationResult:
    confidence: ConfidenceLevel
    score: float
    action: str

class CRAGPipeline:
    """Corrective RAG 风格的流水线实现"""

    def __init__(
        self,
        retriever,
        evaluator,
        generator,
        web_searcher,
        upper_threshold: float = 0.7,
        lower_threshold: float = 0.3
    ):
        self.retriever = retriever
        self.evaluator = evaluator
        self.generator = generator
        self.web_searcher = web_searcher
        self.upper_threshold = upper_threshold
        self.lower_threshold = lower_threshold

    def evaluate_retrieval(self, query: str, documents: list[dict]) -> EvaluationResult:
        """评估检索结果的置信度"""
        scores = []
        for doc in documents:
            score = self.evaluator.score(query, doc["text"])
            scores.append(score)

        max_score = max(scores) if scores else 0.0

        if max_score >= self.upper_threshold:
            return EvaluationResult(
                confidence=ConfidenceLevel.CORRECT,
                score=max_score,
                action="refine_and_use"
            )
        elif max_score <= self.lower_threshold:
            return EvaluationResult(
                confidence=ConfidenceLevel.INCORRECT,
                score=max_score,
                action="web_search_fallback"
            )
        else:
            return EvaluationResult(
                confidence=ConfidenceLevel.AMBIGUOUS,
                score=max_score,
                action="combine_sources"
            )

    def decompose_then_recompose(
        self, query: str, document: str
    ) -> str:
        """Decompose-then-Recompose：仅从文档中提取相关知识"""
        # 第 1 步：将文档分解为细粒度的知识单元
        sentences = document.split(". ")
        knowledge_strips = [s.strip() + "." for s in sentences if s.strip()]

        # 第 2 步：评估每个知识单元的相关性
        relevant_strips = []
        for strip in knowledge_strips:
            relevance = self.evaluator.score(query, strip)
            if relevance > 0.5:
                relevant_strips.append((strip, relevance))

        # 第 3 步：按相关性排序并重新组合
        relevant_strips.sort(key=lambda x: x[1], reverse=True)
        refined_context = " ".join([s[0] for s in relevant_strips])

        return refined_context if refined_context else document

    def process_query(self, query: str) -> dict:
        """执行完整的 CRAG 流水线"""
        # 第 1 步：初始文档检索
        documents = self.retriever.retrieve(query, top_k=10)

        # 第 2 步：评估检索质量
        evaluation = self.evaluate_retrieval(query, documents)

        context = ""
        sources = []

        if evaluation.confidence == ConfidenceLevel.CORRECT:
            # 信任检索结果 —— 提炼核心知识后使用
            for doc in documents[:3]:
                refined = self.decompose_then_recompose(query, doc["text"])
                context += refined + "\n"
            sources = ["internal_retrieval"]

        elif evaluation.confidence == ConfidenceLevel.INCORRECT:
            # 不信任检索结果 —— 转向网络搜索
            web_results = self.web_searcher.search(query)
            for result in web_results[:3]:
                context += result["snippet"] + "\n"
            sources = ["web_search"]

        else:  # AMBIGUOUS
            # 结合两种来源
            for doc in documents[:2]:
                refined = self.decompose_then_recompose(query, doc["text"])
                context += refined + "\n"
            web_results = self.web_searcher.search(query)
            for result in web_results[:2]:
                context += result["snippet"] + "\n"
            sources = ["internal_retrieval", "web_search"]

        # 第 3 步：生成最终回答
        prompt = f"Context: {context}\nQuery: {query}\nAnswer:"
        response = self.generator.generate(prompt)

        return {
            "response": response,
            "confidence": evaluation.confidence.value,
            "score": evaluation.score,
            "sources": sources
        }

import numpy as np
from dataclasses import dataclass, field
from typing import Optional
import re
from collections import Counter
import math

@dataclass
class Document:
    text: str
    doc_id: str
    metadata: dict = field(default_factory=dict)

@dataclass
class SearchResult:
    document: Document
    score: float
    source: str  # "sparse", "dense", or "hybrid"

class BM25Retriever:
    """BM25 Sparse Retriever 的简化实现"""

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.documents: list[Document] = []
        self.doc_freqs: dict[str, int] = {}
        self.doc_lengths: list[int] = []
        self.avg_doc_length: float = 0
        self.doc_term_freqs: list[dict[str, int]] = []

    def _tokenize(self, text: str) -> list[str]:
        return re.findall(r'\w+', text.lower())

    def index(self, documents: list[Document]):
        self.documents = documents
        for doc in documents:
            tokens = self._tokenize(doc.text)
            self.doc_lengths.append(len(tokens))
            term_freq = Counter(tokens)
            self.doc_term_freqs.append(term_freq)
            for term in set(tokens):
                self.doc_freqs[term] = self.doc_freqs.get(term, 0) + 1

        self.avg_doc_length = (
            sum(self.doc_lengths) / len(self.doc_lengths) if self.doc_lengths else 0
        )

    def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
        query_tokens = self._tokenize(query)
        n_docs = len(self.documents)
        scores = []

        for i, doc in enumerate(self.documents):
            score = 0.0
            for term in query_tokens:
                if term not in self.doc_term_freqs[i]:
                    continue
                tf = self.doc_term_freqs[i][term]
                df = self.doc_freqs.get(term, 0)
                idf = math.log((n_docs - df + 0.5) / (df + 0.5) + 1)
                dl = self.doc_lengths[i]
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (
                    1 - self.b + self.b * dl / self.avg_doc_length
                )
                score += idf * numerator / denominator
            scores.append(score)

        top_indices = np.argsort(scores)[::-1][:top_k]
        return [
            SearchResult(
                document=self.documents[i],
                score=float(scores[i]),
                source="sparse"
            )
            for i in top_indices if scores[i] > 0
        ]

class DenseRetriever:
    """Dense Vector Retriever（基于向量嵌入）"""

    def __init__(self, embedding_fn):
        self.embedding_fn = embedding_fn
        self.documents: list[Document] = []
        self.embeddings: Optional[np.ndarray] = None

    def index(self, documents: list[Document]):
        self.documents = documents
        texts = [doc.text for doc in documents]
        self.embeddings = self.embedding_fn(texts)
        # L2 归一化
        norms = np.linalg.norm(self.embeddings, axis=1, keepdims=True)
        self.embeddings = self.embeddings / (norms + 1e-10)

    def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
        q_emb = self.embedding_fn([query])
        q_emb = q_emb / (np.linalg.norm(q_emb) + 1e-10)
        scores = np.dot(self.embeddings, q_emb.T).squeeze()
        top_indices = np.argsort(scores)[::-1][:top_k]
        return [
            SearchResult(
                document=self.documents[i],
                score=float(scores[i]),
                source="dense"
            )
            for i in top_indices
        ]

class HybridRetriever:
    """Hybrid Retrieval：结合 BM25 与 Dense 检索"""

    def __init__(
        self,
        sparse: BM25Retriever,
        dense: DenseRetriever,
        alpha: float = 0.5
    ):
        self.sparse = sparse
        self.dense = dense
        self.alpha = alpha  # Dense 权重（1-alpha 为 Sparse 权重）

    def _normalize_scores(self, results: list[SearchResult]) -> dict[str, float]:
        """Min-Max 归一化"""
        if not results:
            return {}
        scores = [r.score for r in results]
        min_s, max_s = min(scores), max(scores)
        range_s = max_s - min_s if max_s != min_s else 1.0
        return {
            r.document.doc_id: (r.score - min_s) / range_s
            for r in results
        }

    def search(self, query: str, top_k: int = 10) -> list[SearchResult]:
        """基于 Reciprocal Rank Fusion 的混合检索"""
        sparse_results = self.sparse.search(query, top_k=top_k * 2)
        dense_results = self.dense.search(query, top_k=top_k * 2)

        sparse_scores = self._normalize_scores(sparse_results)
        dense_scores = self._normalize_scores(dense_results)

        # 收集所有唯一文档
        all_doc_ids = set(sparse_scores.keys()) | set(dense_scores.keys())
        doc_map = {}
        for r in sparse_results + dense_results:
            doc_map[r.document.doc_id] = r.document

        # 加权组合
        hybrid_scores = {}
        for doc_id in all_doc_ids:
            s_score = sparse_scores.get(doc_id, 0.0)
            d_score = dense_scores.get(doc_id, 0.0)
            hybrid_scores[doc_id] = (
                (1 - self.alpha) * s_score + self.alpha * d_score
            )

        # 排序并返回前 k 个
        sorted_docs = sorted(
            hybrid_scores.items(), key=lambda x: x[1], reverse=True
        )[:top_k]

        return [
            SearchResult(
                document=doc_map[doc_id],
                score=score,
                source="hybrid"
            )
            for doc_id, score in sorted_docs
        ]


# 使用示例
bm25 = BM25Retriever()
docs = [
    Document("RAG는 검색과 생성을 결합한다.", "doc1"),
    Document("Self-RAG는 반성 토큰을 사용한다.", "doc2"),
    Document("CRAG는 검색 품질을 평가한다.", "doc3"),
    Document("RETRO는 2조 토큰 데이터베이스를 사용한다.", "doc4"),
]
bm25.index(docs)
sparse_results = bm25.search("RAG에서 검색 품질 평가 방법은?")
for r in sparse_results:
    print(f"[BM25 Score: {r.score:.4f}] {r.document.text}")

class SemanticChunker:
    """基于语义的分块：用嵌入相似度探测自然边界"""

    def __init__(self, embedding_fn, similarity_threshold: float = 0.75):
        self.embedding_fn = embedding_fn
        self.threshold = similarity_threshold

    def chunk(self, text: str, min_chunk_size: int = 100) -> list[str]:
        """在句子间语义相似度发生变化的地方切分"""
        sentences = [s.strip() for s in text.split(". ") if s.strip()]
        if len(sentences) <= 1:
            return [text]

        # 计算每个句子的嵌入
        embeddings = self.embedding_fn(sentences)

        # 计算相邻句子间的余弦相似度
        chunks = []
        current_chunk = [sentences[0]]

        for i in range(1, len(sentences)):
            sim = np.dot(embeddings[i], embeddings[i - 1]) / (
                np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[i - 1])
                + 1e-10
            )

            if sim < self.threshold and len(". ".join(current_chunk)) >= min_chunk_size:
                # 若相似度低于阈值，则开始新的 chunk
                chunks.append(". ".join(current_chunk) + ".")
                current_chunk = [sentences[i]]
            else:
                current_chunk.append(sentences[i])

        if current_chunk:
            chunks.append(". ".join(current_chunk) + ".")

        return chunks