AI

Python으로 검색 엔진 성능 비교하기: FAISS vs ChromaDB

min민 2024. 12. 26.

SentenceTransformer와 FAISS 및 ChromaDB를 활용한 임베딩 검색 성능 비교

 

이번 글에서는 문장을 벡터(임베딩)로 변환하여 검색하는 두 가지 도구인 FAISSChromaDB를 활용한 검색 성능 비교를 진행하였다. 두 도구를 활용해 검색 속도와 초기 설정 시간 등을 벤치마킹하고, 이를 바탕으로 어떤 상황에서 어떤 도구를 선택해야 하는지 판단할 수 있도록 하였다.

이 글에서는 벤치마크를 진행한 코드의 세부 내용과 각 라이브러리의 특징, 사용 방법을 다룬다.

 

 

주요 코드 설명

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. 앞으로 사용 사례에 따라 두 도구를 적절히 활용하면 검색 성능과 데이터 관리 효율성을 모두 극대화할 수 있다.

 

댓글