Chroma는 가볍고 사용하기 쉬운 벡터 데이터베이스입니다. 본 포스트에서는 Chroma DB를 사용하여 문서를 저장하고, 유사도 검색을 수행하는 방법을 정리합니다.
준비 사항
- Python 3.9 이상
- langchain-chroma 패키지 설치
- 임베딩 모델 준비
라이브러리 설치
Chroma DB와 관련 패키지를 설치합니다.
pip install langchain-chroma langchain-huggingface langchain-community langchain-text-splitters| 패키지 | 설명 |
|---|---|
langchain-chroma | Chroma 벡터 스토어 |
langchain-huggingface | HuggingFace 임베딩 모델 |
langchain-community | TextLoader 등 Document Loaders |
langchain-text-splitters | 문서 분할기 |
문서 준비하기
먼저 벡터 스토어에 저장할 문서를 준비합니다.
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
# 1. 문서 로드
loader = TextLoader("./docs/sample.txt", encoding="utf-8")
# 2. 문서 분할
splitter = RecursiveCharacterTextSplitter(
chunk_size=200,
chunk_overlap=50
)
documents = loader.load_and_split(splitter)
print(f"준비된 청크 수: {len(documents)}")Chroma DB 생성하기
from_documents() 메서드
문서와 임베딩 모델을 전달하여 벡터 스토어를 생성합니다.
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
# 임베딩 모델 준비
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3")
# Chroma DB 생성
db = Chroma.from_documents(
documents=documents,
embedding=embeddings
)
print("Chroma DB 생성 완료!")유사도 검색하기
similarity_search()
질문과 유사한 문서를 검색합니다.
# 유사도 검색
query = "인공지능이란 무엇인가?"
results = db.similarity_search(query, k=3)
print(f"검색 결과: {len(results)}개")
for i, doc in enumerate(results):
print(f"\n[결과 {i+1}]")
print(doc.page_content)주요 파라미터
| 파라미터 | 설명 |
|---|---|
query | 검색할 질문 |
k | 반환할 문서 개수 (기본: 4) |
점수와 함께 검색하기
similarity_search_with_score()
유사도 점수(거리)와 함께 결과를 반환합니다. 점수가 낮을수록 유사합니다.
results = db.similarity_search_with_score("인공지능", k=3)
for doc, score in results:
print(f"점수: {score:.4f}")
print(f"내용: {doc.page_content[:50]}...")
print("---")similarity_search_with_relevance_scores()
0~1 사이의 관련성 점수를 반환합니다. 점수가 높을수록 유사합니다.
results = db.similarity_search_with_relevance_scores("인공지능", k=3)
for doc, score in results:
print(f"관련성: {score:.4f}")
print(f"내용: {doc.page_content[:50]}...")
print("---")Chroma DB 저장하기
persist_directory를 지정하면 DB를 디스크에 저장할 수 있습니다.
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3")
# 저장 경로 지정하여 생성
db = Chroma.from_documents(
documents=documents,
embedding=embeddings,
persist_directory="./chroma_db"
)
print("Chroma DB 저장 완료!")저장된 파일 구조
./chroma_db/
├── chroma.sqlite3
└── [기타 인덱스 파일들]Chroma DB 불러오기
저장된 DB를 다시 불러올 수 있습니다.
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
# 동일한 임베딩 모델 사용
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3")
# 저장된 DB 불러오기
db = Chroma(
persist_directory="./chroma_db",
embedding_function=embeddings
)
# 바로 검색 가능
results = db.similarity_search("인공지능", k=3)
print(f"검색 결과: {len(results)}개")중요: 저장할 때와 불러올 때 반드시 동일한 임베딩 모델을 사용해야 합니다.
저장 vs 불러오기 비교
| 메서드 | 용도 | 사용 시점 |
|---|---|---|
Chroma.from_documents() | 새 DB 생성 및 저장 | 최초 생성 시 |
Chroma() | 기존 DB 불러오기 | 재사용 시 |
전체 예제 코드
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
# 1. 문서 로드 및 분할
loader = TextLoader("./docs/sample.txt", encoding="utf-8")
splitter = RecursiveCharacterTextSplitter(
chunk_size=200,
chunk_overlap=50
)
documents = loader.load_and_split(splitter)
print(f"청크 수: {len(documents)}")
# 2. 임베딩 모델 준비
embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-m3")
# 3. Chroma DB 생성 및 저장
db = Chroma.from_documents(
documents=documents,
embedding=embeddings,
persist_directory="./chroma_db"
)
print("DB 생성 완료!")
# 4. 유사도 검색
query = "인공지능 기술"
results = db.similarity_search_with_relevance_scores(query, k=3)
print(f"\n'{query}' 검색 결과:")
for doc, score in results:
print(f"\n[관련성: {score:.4f}]")
print(doc.page_content[:100])