Python으로 ColBERT RAG 시스템 직접 구축하기: 핵심 모듈별 구현 가이드

Python으로 ColBERT RAG 시스템 직접 구축하기: 핵심 모듈별 구현 가이드

RAG 시스템을 처음 구성하면 대부분 OpenAI의 text-embedding 모델이나 sentence-transformers 기반의 단일 벡터 검색을 쓴다. 빠르게 돌아가고 설정도 간단하다. 그런데 실제로 운영하다 보면 비슷한 문제에 부딪힌다. 쿼리와 문서의 의미는 겹치는데 검색이 빠져버리거나, 단어 하나 차이로 전혀 다른 결과가 나오는 상황이다.

ColBERT는 이 문제를 다른 방식으로 접근한다. 문서 전체를 하나의 벡터로 압축하지 않고, 토큰 단위 임베딩을 그대로 유지한다. 쿼리도 마찬가지다. 그 둘을 나중에 비교(Late Interaction)하는 구조 덕분에, 단일 벡터 방식보다 정밀한 매칭이 가능하다.

이 글에서는 Python으로 ColBERT 기반 RAG 시스템을 직접 구축하는 과정을 핵심 모듈별로 정리한다. RAGatouille 라이브러리를 중심으로, 인덱싱 → 검색 → LangChain 연동 → 도메인 파인튜닝까지 순서대로 다룬다.

---

ColBERT가 일반 임베딩 검색과 다른 이유

단일 벡터(Dense Retrieval) 방식은 문서 하나를 768차원짜리 벡터 하나로 만든다. 쿼리도 벡터 하나로 만든 뒤, 두 벡터의 코사인 유사도를 계산한다. 구조가 단순해서 빠르고 구현하기 쉽다. 하지만 문서 전체의 의미를 벡터 하나에 담으면서 세밀한 정보가 사라진다.

ColBERT는 다르다. 각 토큰마다 임베딩 벡터를 유지한다. "Python으로 파일을 읽는 방법"이라는 쿼리가 있을 때, "Python", "파일", "읽는", "방법" 각각이 벡터를 갖는다. 문서도 마찬가지로 토큰별 벡터를 모두 저장한다.

관련도 점수 계산에는 MaxSim 연산을 쓴다. 각 쿼리 토큰에 대해 문서의 모든 토큰 벡터와 유사도를 계산하고, 그 중 최대값만 취한다. 그 최대값들을 모두 더한 것이 최종 점수다.

Score(Q, D) = Σ max(cos_sim(q_i, d_j) for d_j in D) for q_i in Q

이 방식의 장점은 쿼리 "파일"이라는 토큰이 문서 안에서 가장 관련 있는 위치를 정확히 찾아낸다는 점이다. 문서 전체 의미 평균에 묻혀버리는 단일 벡터와는 결이 다르다.

다만 단점도 있다. 문서당 벡터 수가 토큰 수만큼 늘어나기 때문에 저장 공간이 단일 벡터보다 수십 배 많이 필요하다. 100만 건 문서를 인덱싱하면 디스크 점유가 상당하다. ColBERT v2에서는 벡터 압축(PLAID 엔진)을 도입해 이 문제를 많이 줄였지만, 단일 벡터 방식보다 무겁다는 점은 설계 단계에서 감안해야 한다.

2025년 기준 벤치마크를 보면, 파인낸셜 QA 환경에서 ColBERT 기반 검색은 Recall@5 기준 약 0.816을 기록했다. 같은 환경에서 단순 BM25는 0.644, Dense Retrieval은 0.587이었다. 도메인 특화 문서나 전문 용어가 많은 환경일수록 성능 차이가 두드러진다.

---

환경 설정과 RAGatouille 설치

ColBERT를 직접 사용하는 방법도 있지만, 처음 구축할 때는 RAGatouille를 쓰는 게 현실적이다. Stanford의 ColBERT 공식 구현을 감싸서 인덱싱, 검색, 파인튜닝을 간단한 API로 노출한다.

pip install ragatouille

CUDA 환경이 있으면 GPU 가속이 자동으로 붙는다. CPU 전용 환경에서도 동작하지만 인덱싱 속도가 상당히 느려진다. 수천 건 이상 문서를 다룰 거라면 GPU 환경을 먼저 확인하는 편이 낫다.

설치 후 기본 모델 로드는 이렇게 한다.

from ragatouille import RAGPretrainedModel

# ColBERT v2.0 사전학습 체크포인트 로드
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

colbert-ir/colbertv2.0은 Hugging Face에 올라온 공식 체크포인트다. 영어 기반 제로샷 검색에서는 이 모델만으로도 충분히 쓸 만하다. 한국어 문서라면 jinaai/jina-colbert-v2가 멀티링귀얼 지원이 되어서 시작점으로 쓰기 좋다.

---

인덱싱 파이프라인 구현

문서 컬렉션을 ColBERT 인덱스로 만드는 과정이다. 내부적으로 PLAID 엔진이 토큰 임베딩 압축과 역인덱스 구성을 처리한다.

from ragatouille import RAGPretrainedModel

RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")

# 문서 컬렉션 준비
documents = [
    "Python의 asyncio 모듈은 단일 스레드에서 비동기 I/O를 처리한다.",
    "FastAPI는 Python 타입 힌트를 기반으로 자동 문서를 생성한다.",
    "SQLAlchemy ORM을 사용하면 SQL을 직접 작성하지 않고 DB를 다룰 수 있다.",
    # ... 더 많은 문서
]

document_ids = [f"doc_{i}" for i in range(len(documents))]

# 인덱스 생성
index_path = RAG.index(
    collection=documents,
    document_ids=document_ids,
    index_name="my_python_docs",
    max_document_length=512,  # 문서 최대 토큰 길이
    split_documents=True,     # 긴 문서 자동 분할
)

split_documents=True를 설정하면 max_document_length를 넘는 문서를 자동으로 청크 단위로 나눈다. 직접 청킹 로직을 짜지 않아도 된다는 점에서 편리하다. 단, 분할 기준이 고정 길이 기반이라 문맥이 끊어질 수 있다. 정밀한 청킹이 필요한 경우에는 미리 분할한 문서를 넘기는 방식이 낫다.

메타데이터도 함께 저장할 수 있다.

document_metadatas = [
    {"source": "docs/asyncio.md", "category": "stdlib"},
    {"source": "docs/fastapi.md", "category": "framework"},
    {"source": "docs/sqlalchemy.md", "category": "orm"},
]

RAG.index(
    collection=documents,
    document_ids=document_ids,
    document_metadatas=document_metadatas,
    index_name="my_python_docs",
    max_document_length=512,
)

인덱스는 디스크에 저장된다. 이후에는 모델을 다시 로드할 필요 없이 인덱스만 불러오면 된다.

# 기존 인덱스 불러오기
RAG = RAGPretrainedModel.from_index(".ragatouille/colbert/indexes/my_python_docs")

---

검색 파이프라인 구현

인덱스가 준비되면 검색은 단순하다.

results = RAG.search(
    query="Python에서 비동기 처리 어떻게 하나요?",
    k=5  # 상위 5개 문서 반환
)

for result in results:
    print(f"Score: {result['score']:.4f}")
    print(f"Content: {result['content']}")
    print(f"Metadata: {result['document_metadata']}")
    print("---")

반환되는 결과에는 관련도 점수, 문서 내용, 저장한 메타데이터가 모두 포함된다. 점수는 MaxSim 기반이라 절대값보다는 순위로 해석하는 게 맞다.

복수 쿼리를 한 번에 처리할 수도 있다.

queries = [
    "asyncio 사용법",
    "FastAPI 라우터 설정",
    "SQLAlchemy 세션 관리",
]

batch_results = RAG.search(queries, k=3)
# 각 쿼리별 결과 리스트 반환

검색 결과를 LLM에 넘겨서 답변을 생성하는 흐름은 아래처럼 단순하게 구성할 수 있다.

import anthropic

def rag_query(query: str, rag_model, top_k: int = 3) -> str:
    # 1. ColBERT로 관련 문서 검색
    results = rag_model.search(query=query, k=top_k)

    # 2. 컨텍스트 구성
    context = "\n\n".join([r["content"] for r in results])

    # 3. LLM 호출
    client = anthropic.Anthropic()
    response = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": f"다음 문서를 참고해서 질문에 답해줘.\n\n{context}\n\n질문: {query}"
        }]
    )

    return response.content[0].text

이 방식이 관리하기 편하다. 검색과 생성이 분리되어 있어서 ColBERT 인덱스를 교체하거나 LLM 모델을 바꿔도 나머지 코드에 영향이 없다.

---

LangChain 연동: 리랭커로 활용하기

ColBERT를 완전히 새로운 검색 인덱스로 쓰지 않고, 기존 벡터 스토어 위에 리랭커로 얹는 방식도 자주 쓴다. 이미 운영 중인 Chroma나 Pinecone 인덱스가 있을 때 교체 비용 없이 검색 품질만 높일 수 있다.

LangChain의 ContextualCompressionRetriever를 활용한다.

from ragatouille import RAGPretrainedModel
from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings

# 기존 Dense 검색 리트리버
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(embedding_function=embeddings)
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})

# ColBERT를 리랭커로 설정
RAG = RAGPretrainedModel.from_pretrained("colbert-ir/colbertv2.0")
compressor = RAG.as_langchain_document_compressor(k=5)

# 두 단계 결합: Dense 검색 20개 → ColBERT 리랭킹 → 상위 5개
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever
)

# 사용
docs = compression_retriever.get_relevant_documents("비동기 처리 방법")

이 패턴은 실무에서 생각보다 많이 쓰는 방식이다. Dense 검색으로 후보를 넓게 뽑고 (recall 확보), ColBERT로 정밀하게 순위를 재정렬한다 (precision 개선). 기존 인프라를 건드리지 않아서 도입이 빠르다.

단, 리랭킹 단계에서 ColBERT가 문서를 다시 인코딩해야 해서 레이턴시가 늘어난다. 20개 문서를 리랭킹하면 수백 밀리초 정도 추가된다. 실시간 응답이 중요한 서비스라면 인덱스 방식과 비교해서 결정하는 게 좋다.

---

도메인 특화 파인튜닝

ColBERT v2 제로샷 모델은 일반적인 검색에서는 잘 동작한다. 하지만 도메인 용어가 많거나 특수한 검색 패턴이 있는 경우, 자체 데이터로 파인튜닝하면 성능이 꽤 오른다.

RAGatouille의 RAGTrainer가 이 과정을 처리한다.

from ragatouille import RAGTrainer

trainer = RAGTrainer(
    model_name="my-domain-colbert",
    pretrained_model_name="colbert-ir/colbertv2.0",  # 파인튜닝 기반
    language_code="ko",
)

# 학습 데이터 준비 (쿼리-관련문서-비관련문서 트리플렛)
# RAGatouille는 다양한 형식을 자동으로 변환해준다
training_data = [
    ("Python asyncio 사용법", "asyncio는 단일 스레드 비동기...", "BM25 알고리즘은..."),
    ("FastAPI 라우터 설정", "FastAPI에서 라우터는 APIRouter를...", "Django 설정 파일..."),
]

trainer.prepare_training_data(
    raw_data=training_data,
    data_out_path="./training_data/",
    all_documents=documents,  # 전체 문서 컬렉션 (네거티브 마이닝용)
    mine_hard_negatives=True,  # 어려운 비관련 문서 자동 생성
)

trainer.train(
    batch_size=32,
    nbits=2,       # 압축 비트 수 (저장 공간 vs 품질 트레이드오프)
    maxsteps=500,
    use_ib_negatives=True,  # In-Batch 네거티브 활용
)

mine_hard_negatives=True를 쓰면 단순 랜덤 비관련 문서 대신, 모델이 헷갈릴 만한 어려운 비관련 문서를 자동으로 찾아준다. 이 옵션이 없으면 파인튜닝 효과가 기대보다 낮게 나오는 경우가 많다.

어노테이션 없이 학습하는 방법도 있다. Instructor 같은 LLM을 활용해 합성 데이터(Synthetic Data)를 생성하고 학습에 쓰는 방식인데, 직접 레이블링 비용 없이도 도메인 적응이 가능하다. 실제 시도해보면 어노테이션 데이터와 비교해서 크게 뒤지지 않는 경우가 많다.

---

실전 팁과 자주 막히는 부분

인덱싱이 멈추거나 OOM이 발생할 때

max_document_length를 낮추는 것이 빠른 해결책이다. 기본값이 너무 크면 GPU 메모리를 초과할 수 있다. 256 또는 128로 낮춰서 시작해보는 걸 권장한다. 처리 속도와 품질 사이에서 조율이 필요하다.

한국어 문서에서 검색 품질이 낮을 때

colbert-ir/colbertv2.0은 영어 기반 모델이다. 한국어 문서를 다룰 때는 jinaai/jina-colbert-v2(멀티링귀얼) 또는 한국어 사전학습 BERT 기반으로 파인튜닝한 체크포인트를 써야 제대로 된 성능이 나온다. 영어 모델로 한국어 문서를 인덱싱하면 토큰화 자체에서 정보 손실이 생긴다.

기존 Chroma/Pinecone 데이터와 병행 운영할 때

리랭커 방식(위의 LangChain 연동)을 먼저 적용해 효과를 확인하고, 충분히 검증되면 전체 인덱스를 ColBERT로 마이그레이션하는 순서가 안전하다. 한 번에 바꾸려고 하면 기존 운영 중인 서비스에 영향을 주기 쉽다.

PLAID 인덱스와 일반 인덱스 차이

RAGatouille 내부에서 ColBERT v2는 기본적으로 PLAID 엔진을 사용한다. 압축된 토큰 임베딩과 역인덱스 구조를 결합해서 대규모 코퍼스에서도 낮은 레이턴시를 유지한다. 1억 4천만 개 패시지 규모에서도 서브세컨드 응답이 가능하다는 논문 결과가 있다. 직접 100만 건 이하 규모로 테스트해보면 단일 벡터 검색과 응답 속도 차이가 체감될 만큼 크지 않다.

비용 관련

ColBERT는 자체 호스팅 기반이다. OpenAI Embedding API처럼 건당 비용이 발생하지 않는다. GPU 서버 비용은 들지만, 대용량 문서를 자주 검색하는 경우 장기적으로 API 비용보다 저렴해질 수 있다.

---

마무리

ColBERT RAG 시스템 구축을 정리하면 이렇다.

• 기본 설치: pip install ragatouille 한 줄
• 인덱싱: RAG.index()로 문서 컬렉션 → PLAID 인덱스 생성
• 검색: RAG.search()로 MaxSim 기반 관련 문서 반환
• LangChain 연동: as_langchain_document_compressor()로 리랭커로 활용
• 파인튜닝: RAGTrainer로 도메인 특화 성능 향상

단일 벡터 검색으로 커버가 안 되는 케이스가 반복적으로 발생하거나, 도메인 전문 문서를 다루는 시스템이라면 ColBERT를 검토해볼 시점이다. 설치와 인덱싱만 해보는 데는 30분이면 충분하다. 기존 벡터 스토어 위에 리랭커로 얹는 방식부터 시작하면 기존 코드 변경 없이 효과를 먼저 확인할 수 있다.

---

Sources:
- GitHub - AnswerDotAI/RAGatouille
- RAGatouille · PyPI
- RAGatouille | LangChain
- Ragatouille integrations - Docs by LangChain
- Introduction to ColBERT | RAGStack | DataStax Docs
- Exploring ColBERT with RAGatouille | Simon Willison's TILs
- colbert-ai · PyPI
- GitHub - stanford-futuredata/ColBERT
- PLAID: An Efficient Engine for Late Interaction Retrieval
- ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interaction
- jinaai/jina-colbert-v2 · Hugging Face
- ColBERT: A Token-Level Embedding and Ranking Model - Zilliz Learn
- An Overview of Late Interaction Retrieval Models: ColBERT, ColPali, and ColQwen | Weaviate
- ColBERT in Practice: Bridging Research and Industry - Sease.io
- Production Retrivers in RAG That Work: Hybrid Search + Re-Ranking | Medium