오픈 소스를 활용한 RAG 구현

RAG를 구성하는 쿼리 임베딩 모델, 문서 임베딩 모델, 생성 모델을 오픈 소스 모델로 구성해보았습니다. 서버가 갖춰진 경우 추가로 비용이 발생한다는 장점이 있지만, 서버가 구비되어 있는 경우가 흔치는 않을 것 같습니다. 과제를 수행하는 환경 때문에 OpenAI의 text-embedding-ada-002 모델을 사용할 수 없어 시도해본 방법입니다.

• 서버 환경은 도커 이미지(pytorch/pytorch:2.1.1-cuda12.1-cudnn8-runtime)를 이용해 구성했습니다.
• GPU는 NVIDIA GeForce RTX 3090 Ti 2장 이용했습니다.
• CUDA Driver Version은 12.0입니다.
• 임베딩 모델은 JinaAI가 개발한 오픈소스 임베딩 모델을 사용했습니다.
• 생성 모델은 LLAMA2를 이용했습니다.

아래 페이지의 글을 참고했습니다. 본 글의 코드와 내용은 아래 페이지의 간략한 버전이라고 생각해주세요.

• https://blog.devgenius.io/build-your-first-conversational-document-retrieval-agent-using-llama-2-and-langchain-a7adf5e50a63
• https://rfriend.tistory.com/826 [R, Python 분석과 프로그래밍의 친구 (by R Friend):티스토리]

---

하나의 문서를 참고해 답변하는 RAG를 만들어보겠습니다.

documents = ["harrison worked at kensho"]
query = "Where did harrison work?"

필요한 라이브러리를 설치합니다.

!pip install -q accelerate==0.21.0 bitsandbytes==0.40.2 langchain tiktoken faiss-cpu sentence-transformers

필요한 라이브러리를 임포트합니다.

from langchain import HuggingFacePipeline
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain.vectorstores import FAISS
from langchain.chains import ConversationalRetrievalChain

import torch
from transformers import pipeline, AutoTokenizer

Embedding 모델을 로드합니다.

# Load Embedding
# Define the path to the pre-trained model you want to use
modelPath = "jinaai/jina-embeddings-v2-small-en" # open source embedding model

# Create a dictionary with model configuration options, specifying to use the CPU for computations
model_kwargs = {'device': 'cpu'}

# Create a dictionary with encoding options, specifically setting 'normalize_embeddings' to False
encode_kwargs = {'normalize_embeddings': False}

# Initialize an instance of HuggingFaceEmbeddings with the specified parameters
embeddings = HuggingFaceEmbeddings(
    model_name=modelPath,     # Provide the pre-trained model's path
    model_kwargs=model_kwargs, # Pass the model configuration options
    encode_kwargs=encode_kwargs # Pass the encoding options
)

벡터 DB를 만듭니다.

# Create a FAISS VectorStore
vectorstore = FAISS.from_texts(
    documents, embedding=embeddings
)

# Create a Retriever
retriever = vectorstore.as_retriever()

프롬프트 템플릿을 만듭니다.

# Create Template
template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

라마 모델과 토크나이저를 로드하고, 바로 문자열을 뱉도록 LLM 파이프라인을 만듭니다.

# Create a tokenizer, pipeline
model_name = "NousResearch/llama-2-7b-chat-hf"

# Load LLaMA tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right" # Fix weird overflow issue with fp16 training

pipe = pipeline(task="text-generation",
                model=model_name,
                tokenizer=tokenizer,
                return_full_text=True,
                max_length=400,
                device_map="auto"
               )

llm = HuggingFacePipeline(pipeline = pipe)

RAG 파이프라인을 만듭니다.

# Chaining a pipeline
chain = ConversationalRetrievalChain.from_llm(
    llm,
    retriever,
    return_source_documents=True
)

RAG 파이프라인에 질문을 던져 답변을 확인합니다. chat_history를 이용해 multi-turn 대화를 이어나갈 수 있습니다.

# Run retrieva_chain
chat_history = []
result = chain({"question": query, "chat_history": chat_history})

print(result['answer'])
# Expected Answer : 'Harrison worked at Kensho.'