1. 기본 설정 및 컬렉션 생성
import chromadb
from chromadb.config import Settings
# ChromaDB 클라이언트 생성
client = chromadb.Client(Settings(
chroma_db_impl="duckdb+parquet",
persist_directory="./chroma_storage" # 데이터 영구 저장 경로
))
# 컬렉션 생성
collection = client.create_collection(
name="book_collection",
metadata={"description": "책 정보 저장소"}
)
2. 데이터 임베딩 및 추가
from sentence_transformers import SentenceTransformer
# 임베딩 모델 로드
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# 샘플 책 데이터
books = [
"데미안 - 헤르만 헤세의 소설",
"1984 - 조지 오웰의 디스토피아 소설",
"홍길동전 - 허균의 고전 소설"
]
# 임베딩 생성
embeddings = embedding_model.encode(books)
# ChromaDB에 데이터 추가
collection.add(
embeddings=embeddings.tolist(),
documents=books,
ids=[f"book_{i}" for i in range(len(books))]
)
3. 유사도 검색
# 쿼리 임베딩 생성
query = "소설 추천해줘"
query_embedding = embedding_model.encode([query]).tolist()
# 유사도 검색 수행
results = collection.query(
query_embeddings=query_embedding,
n_results=2 # 상위 2개 결과 반환
)
# 결과 출력
for i, doc in enumerate(results['documents'][0]):
print(f"유사도 결과 {i+1}: {doc}")
print(f"유사도 점수: {results['distances'][0][i]}")
4. 고급 필터링
# 메타데이터를 포함한 컬렉션 생성
collection_with_meta = client.create_collection(
name="advanced_book_collection",
metadata={"type": "book_catalog"}
)
# 메타데이터 포함 데이터 추가
collection_with_meta.add(
embeddings=embeddings.tolist(),
documents=books,
ids=[f"book_{i}" for i in range(len(books))],
metadatas=[
{"genre": "소설", "author": "헤르만 헤세"},
{"genre": "디스토피아", "author": "조지 오웰"},
{"genre": "고전", "author": "허균"}
]
)
# 메타데이터 기반 필터링
filtered_results = collection_with_meta.query(
query_embeddings=query_embedding,
where={"genre": "소설"}, # 소설 장르만 필터링
n_results=1
)
5. 컬렉션 관리
# 컬렉션 목록 확인
print(client.list_collections())
# 특정 컬렉션 삭제
client.delete_collection(name="book_collection")
