[Feat] Retrieval Augmented Generation 시스템 구축

[eyeol]

📝 Summary

• Pinecone 기반의 Vector DB와 연결된 RAG 시스템을 구현했습니다.
• 앞선 실험들의 결과를 참고하여 고품질의 Vector DB를 구성할만한 문서를 선별했습니다.
• 직접 개발한 쿼리 빌더 클래스를 사용하여 다양한 쿼리에 대한 실험을 진행하며, LLM as Judege 방식으로 retrieve된 문서를 평가하였습니다.
• 모델의 견고성을 높이기 위해 RAFT 방식으로 학습할 수 있는 환경을 구축했습니다.

✅ Checklist

• 관련 이슈가 명시되어 있습니다.
• 테스트가 완료되었습니다.
• 문서 업데이트가 포함되었습니다.
• 코드 리뷰를 위한 사전 검토를 완료했습니다.

📄 Description

RAG 시스템 구축

RAG (Retrieval-Augmented Generation) 시스템을 구현하여, 질의 응답 정확도를 높이고 외부 데이터와의 연결을 강화하였습니다. 이번 작업에서는 다음과 같은 체인을 구성하였습니다:

1. Pinecone 기반 Retriever

# rag_modules/retriever.py
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_pinecone import PineconeVectorStore

def get_retriever(embedding_model: str = "BAAI/bge-m3", k: int = 5, minimal_score: float = 0.4):
    # Pinecone index 연결
    pinecone_index, index_name = get_pinecone_index()

    # HuggingFace Embeddings 모델 초기화
    embeddings = HuggingFaceBgeEmbeddings(model_name=embedding_model)

    # LangChain의 Pinecone VectorStore 생성
    vector_store = PineconeVectorStore(index=pinecone_index, embedding=embeddings, text_key="text")

    # Retriever 반환
    return vector_store.as_retriever(
        search_type="similarity_score_threshold", search_kwargs={"k": k, "score_threshold": minimal_score}
    )

Pinecone의 Vector Store 객체를 활용하여 Dense Retriever를 구현하였습니다. as_retriever() 메서드를 사용해 Pinecone Vector Store를 간단하게 검색 인터페이스로 변환했으며, 유사도 검색에는 Pinecone에 설정된 metric (예: cosine, dot_product 등)이 사용됩니다.

2. 생성(추론) 모델 chain

def create_chain(model_id: str = "google/gemma-2-2b-it", max_new_tokens: int = 256):
    # Initialize HuggingFace LLM
    model = AutoModelForCausalLM.from_pretrained(model_id)
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    gen = pipeline(
        task="text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=max_new_tokens,
        device=0 if torch.cuda.is_available() else -1,
        return_full_text=False,
    )
    llm = HuggingFacePipeline(pipeline=gen)

    parser = JsonOutputParser(pydantic_object=Output)
    fixing_parser = OutputFixingParser.from_llm(parser=parser, llm=llm)

    chat = [{"role": "user", "content": template}]

    prompt_template = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)

    # Create prompt template
    prompt = PromptTemplate.from_template(
        template=prompt_template, partial_variables={"format": parser.get_format_instructions()}
    )

    # Combine into chain
    return prompt | llm | fixing_parser

이 체인은 RAG 시스템에서 retriever의 결과물을 받아 질의 응답의 최종 단계를 처리합니다. Prompt Template, LLM, 그리고 Output Parser를 결합하여 검색된 데이터와 함께 모델의 능력을 활용해 구조화된 응답을 생성합니다.

Prompt Template은 검색된 문서를 컨텍스트로 포함하며, HuggingFace LLM은 텍스트 생성을 수행합니다. 마지막으로 Output Parser는 생성된 텍스트를 구조화된 형식으로 변환하여, 최종 응답으로 반환합니다. 이 체인은 검색과 생성 단계를 연결하며, RAG 시스템의 질의 응답 정확도와 일관성을 보장하는 핵심적인 역할을 합니다.

Pinecone 이용한 Vector DB

1. Pinecone을 선택한 이유

Pinecone은 클라우드 기반 벡터 데이터베이스로, 협업 시 관리와 활용이 용이합니다.
FAISS와 같은 로컬 솔루션에 비해 별도의 서버 설정이나 운영 부담 없이 사용할 수 있어 팀 단위 작업에 적합합니다.
또한, 앞서 작성한 retriever 코드에서 볼 수 있듯이, 간단한 API와 개발 편의성을 제공해 벡터 데이터베이스를 쉽게 통합할 수 있다는 장점이 있습니다.

2. Vector DB에 넣을 문서의 선정 기준

모든 위키피디아 문서를 벡터 DB에 넣을 경우, 검색된 문서의 품질이 떨어질 위험이 있다고 판단했습니다. 이를 방지하기 위해, 앞선 실험에서 도출된 키워드(자세한 내용은 별도로 기술)를 중심으로 관련된 문서를 크롤링하여 데이터셋을 구성하였습니다. 크롤링한 문서들은 HuggingFace 임베딩 모델을 사용해 벡터화한 뒤 Pinecone에 저장했습니다.

3. 문서 임베딩 후 VectorDB에 넣는 과정

아래 코드는 문서를 처리하여 Pinecone에 업로드하는 과정을 보여줍니다. 문서를 청크 단위로 나누어 임베딩을 생성하고, 각 청크를 고유 ID와 메타데이터와 함께 벡터 DB에 저장합니다:

def rag_preprocessing(page_text_path, page_title, page_index):
    # Load and process document
    documents = load_document(page_text_path)
    chunks = split_into_chunks(documents)
    # print(chunks)
    embeddings = generate_embeddings(chunks, embedding_type="huggingface")

    # Prepare and upsert embeddings with unique IDs
    upsert_data = []
    for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
        unique_id = f"{chunk.metadata.get('document_id', 'unknown')}_{i}_{uuid.uuid4()}"

        upsert_data.append(
            {
                "id": unique_id,  # 고유한 ID 생성
                "values": embedding,
                "metadata": {
                    "text": chunk.page_content,
                    "document_id": page_title,
                    "chunk_index": i,
                    "page_index": page_index,
                },
            }
        )

    # Upload embeddings
    pinecone_index.upsert(vectors=upsert_data)

사진 출처

위 이미지는 문서 데이터를 처리하여 Pinecone에 저장하기까지의 전체 흐름을 시각화한 것입니다.

• Load : JSON 파일이나 URL 등 다양한 소스에서 문서를 로드합니다.
• Split : 긴 문서를 작은 청크로 나눠 벡터화하기 적합한 형태로 만듭니다.
• Embed : HuggingFace 임베딩 모델을 사용해 각 청크를 벡터로 변환합니다.
• Store : 생성된 벡터와 메타데이터를 Pinecone의 Vector DB에 저장합니다.

이 흐름은 효율적인 데이터 검색과 모델의 성능을 위한 기반을 제공합니다.

최적의 query 탐색

RAG 시스템의 성능을 높이기 위해 query를 다양한 방식으로 변형하여 실험을 진행했습니다. 초기에는 단순히 paragraph만을 query로 사용했지만, retrieve된 문서가 실제 문제 풀이에 적합하지 않을 가능성을 발견했습니다. 특히, public 데이터셋에서 성능 향상이 나타나지 않으면서 이런 의문이 더욱 강해졌습니다(다만, final 단계에서는 성능이 일부 향상됨).

쿼리 빌더를 활용한 다양한 실험

이에 따라 query에 데이터를 추가하는 다양한 조합을 시도했습니다. 예를 들어, paragraph 뿐만 아니라 데이터 행에서 추출한 paragraph, 그리고 paragraph + question 조합 등을 활용했습니다. 이러한 실험을 더 효율적으로 진행하기 위해, 다양한 쿼리 빌더를 Config로 설정해 손쉽게 변경 및 테스트할 수 있는 환경을 구성하였습니다.

match query_builder_type:
        case "OriginalKeywordsQueryBuilder":
            return OriginalKeywordsQueryBuilder()
        ...

        case "CombinedKeyQueryBuilder_s":
            return CombinedKeyQueryBuilder(["summarization"])
        case "TFIDFQueryBuilder":
            return TFIDFQueryBuilder()

LLM as Judge

쿼리 빌더 실험 중, retrieve된 문서의 품질 평가가 어렵다는 문제를 인지했습니다. retrieve된 문서가 실제로 문제 풀이에 유용한지를 판단할 체계적인 방법이 필요했으며, 이를 위해 LLM을 평가 도구(Judge)로 활용하는 방식을 도입했습니다.

LLM은 retrieve된 문서와 query를 기반으로, 해당 문서가 문제 풀이에 얼마나 유용한지를 점수화하거나, 정성적으로 평가하도록 설계되었습니다. 이를 통해 실험 결과를 보다 명확히 해석할 수 있었으며, query 빌더의 개선 방향을 구체화할 수 있었습니다.

RAFT 방식으로 모델 학습

LLM as Judge로 retrieve된 문서를 평가하면서 동시에, 모델의 robustness를 높이기 위해
RAFT 방식으로 추론 모델을 학습할 수 있는 코드를 구현했습니다.

자세한 내용은 #71 에서 확인해주세요.

💡 Notice (Optional)

🔗 Related Issue(s)

#40 #54 #65

[canolayoo78]

늦은 시간까지 수고 많으셨습니다!
꼼꼼한 정리 감사드립니다.