Decide Then Retrieve: Uncertainty 기반 선택적 검색과 Dual-Path Retrieval을 활용한 Training-Free RAG 프레임워크

import numpy as np

def compute_uncertainty(log_probs: list[float]) -> float:
    """토큰별 log probability 리스트로 불확실도 계산"""
    T = len(log_probs)
    return -sum(log_probs) / T

def should_retrieve(query: str, llm, threshold: float = 0.005) -> tuple[bool, str]:
    """UGT: 불확실도 기반 검색 트리거 판단"""
    result = llm.generate_with_logprobs(query + "\nAnswer the question using a single word or phrase.")
    uncertainty = compute_uncertainty(result.log_probs)
    return uncertainty > threshold, result.text

def dual_path_retrieval(query: str, pseudo_context: str, retriever, n: int = 5) -> list:
    """DPR-AIS: 쿼리 + 가상 문서 두 경로로 검색 후 재점수화"""
    q_emb = retriever.encode(query)
    p_emb = retriever.encode(pseudo_context)
    
    docs_q = retriever.search(q_emb, top_k=n)  # 쿼리 경로
    docs_p = retriever.search(p_emb, top_k=n)  # 가상 문서 경로
    candidates = list(set(docs_q + docs_p))     # 합집합 2n개
    
    # AIS: cos(θ1 + θ2) = s1*s2 - sqrt(1-s1²)*sqrt(1-s2²)
    scored = []
    for doc in candidates:
        d_emb = retriever.encode(doc)
        s1 = np.dot(q_emb, d_emb)  # 쿼리-문서 유사도
        s2 = np.dot(p_emb, d_emb)  # 가상문서-문서 유사도
        joint_score = s1 * s2 - np.sqrt(max(0, 1 - s1**2)) * np.sqrt(max(0, 1 - s2**2))
        scored.append((doc, joint_score))
    
    return [doc for doc, _ in sorted(scored, key=lambda x: -x[1])[:3]]

# 사용 예시
def dtr_answer(query: str, llm, retriever) -> str:
    needs_retrieval, initial_answer = should_retrieve(query, llm, threshold=0.005)
    
    if not needs_retrieval:
        return initial_answer  # 확신 있으면 바로 반환
    
    # 가상 문서 생성 후 dual-path 검색
    pseudo = llm.generate(query + "\nWrite a passage to answer this question.")
    docs = dual_path_retrieval(query, pseudo, retriever)
    
    context = "\n".join(docs)
    return llm.generate(f"{query}\n{context}\nAnswer using a single word or phrase.")