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

        # 시퀀스를 청크 단위로 분할
        h_chunks = hidden_states[:, :n_chunks * self.chunk_size].reshape(
            B, n_chunks, self.chunk_size, D
        )

        # 각 청크에 대해 검색된 이웃과 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)

        # Multi-head 분리
        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)

        # 나머지 토큰 복원(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 토큰 시뮬레이션)"""
        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 토큰 시뮬레이션)"""
        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: 문서에서 관련 지식만 추출"""
        # Step 1: 문서를 세밀한 지식 단위로 분해
        sentences = document.split(". ")
        knowledge_strips = [s.strip() + "." for s in sentences if s.strip()]

        # Step 2: 각 지식 단위의 관련성 평가
        relevant_strips = []
        for strip in knowledge_strips:
            relevance = self.evaluator.score(query, strip)
            if relevance > 0.5:
                relevant_strips.append((strip, relevance))

        # Step 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:
                # 유사도가 임계값 미만이면 새로운 청크 시작
                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