SentenceTransformer와 FAISS 및 ChromaDB를 활용한 임베딩 검색 성능 비교
이번 글에서는 문장을 벡터(임베딩)로 변환하여 검색하는 두 가지 도구인 FAISS와 ChromaDB를 활용한 검색 성능 비교를 진행하였다. 두 도구를 활용해 검색 속도와 초기 설정 시간 등을 벤치마킹하고, 이를 바탕으로 어떤 상황에서 어떤 도구를 선택해야 하는지 판단할 수 있도록 하였다.
이 글에서는 벤치마크를 진행한 코드의 세부 내용과 각 라이브러리의 특징, 사용 방법을 다룬다.
주요 코드 설명
1. 임베딩 준비
문장을 벡터로 변환하기 위해 SentenceTransformer 모델을 사용하였다.
from sentence_transformers import SentenceTransformer
class SearchBenchmark:
def __init__(self):
self.embedder = SentenceTransformer("all-MiniLM-L6-v2")
self.df = pd.read_csv("ChatbotData.csv")
self.setup_data()
def setup_data(self):
# 데이터 준비
self.docs = self.df["Q"].tolist()
self.embeddings = self.embedder.encode(self.docs, normalize_embeddings=True)
- SentenceTransformer: 텍스트 데이터를 고정된 크기의 벡터로 변환하기 위해 사용하였다. all-MiniLM-L6-v2 모델은 효율성과 성능이 적절히 균형을 이루는 모델이다.
- 데이터 준비: 데이터프레임에서 질문(Q) 열을 추출해 리스트로 변환하고, encode 메서드를 통해 임베딩을 생성하였다. 여기서 normalize_embeddings=True 옵션을 사용해 벡터를 정규화(normalize)하였다.
2. FAISS 설정
FAISS는 Facebook에서 개발한 벡터 검색 라이브러리로, 빠른 검색 속도가 특징이다.
@pysnooper.snoop()
def setup_faiss(self):
start_time = time.time()
# FAISS 인덱스 생성
dimension = self.embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(self.embeddings)
end_time = time.time()
return index, end_time - start_time
- IndexFlatL2: L2(유클리드 거리)를 기준으로 벡터 간의 유사도를 계산한다.
- add: 생성된 인덱스에 임베딩 데이터를 추가하였다.
FAISS는 설정이 간단하며, 메모리 상에서만 동작하므로 속도가 매우 빠르다. 하지만 데이터 크기가 메모리에 제한되지 않을 정도로 큰 경우 사용이 어려울 수 있다.
3. ChromaDB 설정
ChromaDB는 벡터 검색을 위한 데이터베이스로, 영속적으로 데이터를 저장할 수 있는 점이 장점이다.
@pysnooper.snoop()
def setup_chroma(self):
start_time = time.time()
# Chroma 설정
client = chromadb.PersistentClient()
collection = client.create_collection('benchmark_test')
# 배치 추가
BATCH_SIZE = 5000
for i in range(0, len(self.docs), BATCH_SIZE):
batch_end = min(i + BATCH_SIZE, len(self.docs))
collection.add(
embeddings=self.embeddings[i:batch_end].tolist(),
documents=self.docs[i:batch_end],
ids=[str(x) for x in range(i, batch_end)]
)
end_time = time.time()
return collection, end_time - start_time
- PersistentClient: 데이터를 영속적으로 저장할 수 있는 클라이언트를 생성하였다.
- create_collection: 새로운 컬렉션을 생성하였다.
- add: 데이터를 배치 단위로 추가하여 처리 속도를 개선하였다.
ChromaDB는 대량의 데이터를 다룰 때 적합하며, 데이터베이스 형태로 저장되므로 데이터를 쉽게 관리하고 수정할 수 있다.
4. 검색 벤치마킹
FAISS와 ChromaDB를 각각 활용하여 동일한 쿼리를 검색한 뒤, 초기 설정 시간과 검색 속도를 비교하였다.
@pysnooper.snoop()
def benchmark_search(self, query="오늘 날씨 어때?", n_iterations=100):
query_embedding = self.embedder.encode(query, normalize_embeddings=True)
# FAISS 검색
index, faiss_setup_time = self.setup_faiss()
faiss_start = time.time()
for _ in range(n_iterations):
distances, indices = index.search(np.expand_dims(query_embedding, axis=0), 10)
faiss_time = (time.time() - faiss_start) / n_iterations
# Chroma 검색
collection, chroma_setup_time = self.setup_chroma()
chroma_start = time.time()
for _ in range(n_iterations):
results = collection.query(
query_embeddings=[query_embedding.tolist()],
n_results=10
)
chroma_time = (time.time() - chroma_start) / n_iterations
return {
'faiss_setup_time': faiss_setup_time,
'chroma_setup_time': chroma_setup_time,
'faiss_search_time': faiss_time,
'chroma_search_time': chroma_time
}
- 동일한 쿼리(예: "오늘 날씨 어때?")에 대해 FAISS와 ChromaDB의 평균 검색 시간을 계산하였다.
- n_iterations=100으로 설정해 여러 번 반복 실행한 평균 시간을 구하였다.
벤치마킹 결과
# 벤치마크 실행
benchmark = SearchBenchmark()
results = benchmark.benchmark_search()
print("\n=== 벤치마크 결과 ===")
print(f"FAISS 초기 설정 시간: {results['faiss_setup_time']:.4f}초")
print(f"Chroma 초기 설정 시간: {results['chroma_setup_time']:.4f}초")
print(f"FAISS 평균 검색 시간: {results['faiss_search_time']*1000:.4f}ms")
print(f"Chroma 평균 검색 시간: {results['chroma_search_time']*1000:.4f}ms")
결과를 통해 FAISS는 초기 설정 시간이 짧고 검색 속도가 빠르지만, 데이터의 영속적인 저장이 필요하지 않은 경우에 적합하다는 결론을 얻었다. 반면, ChromaDB는 설정 시간이 오래 걸리지만 데이터를 영속적으로 관리할 수 있으므로 대량의 데이터와 장기적인 사용이 필요한 경우에 적합하였다.
총 정리
1. FAISS와 ChromaDB는 각각의 장단점을 가지고 있다.
2. 검색 속도가 중요하고 메모리 내 데이터로 충분히 처리할 수 있다면 FAISS를 사용하는 것이 적합하다.
3. 반면, 데이터베이스처럼 데이터를 저장하고 관리해야 하는 경우 ChromaDB를 사용하는 것이 더 유리하다.
4. 앞으로 사용 사례에 따라 두 도구를 적절히 활용하면 검색 성능과 데이터 관리 효율성을 모두 극대화할 수 있다.
'AI' 카테고리의 다른 글
Face Recognition 라이브러리와 Google Colab을 활용한 얼굴 인식 구현하기 (0) | 2024.12.27 |
---|
댓글