PostgreSQL과 pgvector를 사용하여 RAG를 구축하는 방법에 대하여 설명합니다. pgvector는 벡터 간 유사도 검색을 지원하는 PostgreSQL extension입니다.
PostgreSQL과 pgvector를 직접 설치할 수도 있지만 여기에서는 미리 구성된 ankane/pgvector 도커 이미지를 사용하여 컨테이너로 구동하겠습니다.
먼저 아래와 같이 도커 이미지를 가져옵니다.
docker pull ankane/pgvector
아래와 같이 컨테이너를 실행합니다. 데이터베이스 사용자와 암호, 데이터베이스 이름 그리고 컨테이너의 이름을 정해줍니다. 컨테이너의 내부 포트는 5432로 고정이지만 외부 포트는 바꾸어도 상관없습니다.
docker run -e POSTGRES_USER=myuser \
-e POSTGRES_PASSWORD=mypassword \
-e POSTGRES_DB=mydatabase \
--name my_postgres \
-p 5432:5432 \
-d ankane/pgvector
컨테이너가 실행되었으면 아래와 같이 psql 툴을 사용하여 데이터베이스에 접속합니다. 여기에서는 로컬 머신에서 컨테이너가 실행중이고 외부 포트가 5432로 매핑된 경우를 가정합니다.
psql -h localhost -U myuser -d mydatabase -p 5432
아래와 같이 pgvector extension을 활성화 해주어야 합니다. 이 작업은 처음 한 번만 수행하면 됩니다.
CREATE EXTENSION vector;
pgvector extension이 활성화된 것을 확인하려면 아래의 SQL 문을 수행하여 vector가 결과에 보이는지 확인하면 됩니다.
SELECT * FROM pg_extension;
파이썬 코드에서 pgvector 데이터베이스에 접근할 때에는 psycopg2를 사용할 수 있습니다.
pip install psycopg2-binary
이제 파이썬 코드로 RAG 테이블을 생성해 보겠습니다. 먼저 아래와 같이 데이터베이스 연결을 생성합니다.
import psycopg2
conn = psycopg2.connect(
host="localhost"
port="5432"
user="myuser",
password="mypassword",
dbname="mydatabase",
)
아래는 RAG라는 이름의 테이블을 생성하는 파이썬 코드입니다. 벡터의 차원은 어떤 임베딩 모델을 사용하는지에 따라 달라집니다. 여기에서는 HuggingFace의 sentence-transformers/all-mpnet-base-v2를 사용하겠습니다. 이 모델은 768차원의 벡터를 생성합니다.
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE RAG (
id SERIAL PRIMARY KEY,
text TEXT,
embedding vector(768)
);
""")
conn.commit()
테이블을 생성했으면 이제 데이터를 로드할 차례입니다. PDF 파일을 로드하여 테이블에 적재해 보겠습니다. 텍스트를 분할하는 청크의 크기는 500으로, 청크 간 텍스트 중첩은 50으로 설정했습니다.
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
# PDF 파일 로드
loader = PyMuPDFLoader("mydata.pdf")
pages = loader.load()
# 텍스트 분할
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50
)
docs = text_splitter.split_documents(pages)
# 벡터 임베딩 생성
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
embeddings = embedding_model.embed_documents([doc.page_content for doc in docs])
# RAG 테이블에 적재
for text, embedding in zip([doc.page_content for doc in docs], embeddings):
cur.execute(
"INSERT INTO RAG (text, embedding) VALUES (%s, %s)",
(text, embedding)
)
conn.commit()
이렇게 도큐먼트가 준비되었으면 아래와 같이 쿼리와 유사한 컨텍스트를 찾아내는 파이썬 함수를 만들어볼 수 있습니다.
def retrieve_context(query, top_k=5):
query_embedding = embedding_model.embed_query(query)
with conn.cursor() as cur:
cur.execute("""
SELECT text, embedding
FROM RAG
ORDER BY embedding <-> %s::real[]::vector(768)
LIMIT %s
""", (query_embedding, top_k))
results = cur.fetchall()
return [row[0] for row in results]
이 함수를 사용하여 쿼리와 유사한 컨텍스트를 찾아서 프롬프트에 추가해주면 LLM으로부터 좀 더 정확한 답을 얻을 수 있습니다. 위에서는 벡터 간 유사도를 비교하는 함수로 <->(L2 distance)를 사용했는데 <#>(inner product), <=>(cosine distance), <+>(L1 distance) 등의 함수도 사용할 수 있습니다.