faiss 사용법 및 예제 코드

FAISS (Facebook AI Similarity Search)의 사용법 및 예제 코드입니다. FAISS는 Facebook AI에서 개발한 라이브러리로 대규모 데이터셋에 대한 빠른 유사도 검색을 할 수 있게 해줍니다.

 

2023. 10. 03 최초작성

from sentence_transformers import SentenceTransformer import pandas as pd import numpy as np import faiss df = pd.DataFrame([[0, 'I like a apple.'], [1, 'I like a orange.'], [2, 'I like a banana.'], [3, 'I like a pineapple.'], [4, 'I like a melon.']], columns=['id', 'text']) # 데이터 프레임은 shape가 (5, 2)입니다. 5개의 인덱스와 5개의 문자열로 구성됩니다. # print(f'{df.shape}') print(df) print() # 문자열을 임베딩 벡터로 변환하기 위해 Sentence Transformer를 사용합니다. model = SentenceTransformer('distilbert-base-nli-stsb-mean-tokens') # 데이터프레임 df의 text 컬럼에 있는 문자열을 인코딩합니다. texts = df.text.values.tolist() embeddings = model.encode(texts) # 5개의 문자열이 저장된 리스트는 shape가 (5, 768)인 벡터로 변환됩니다. # print(embeddings.shape) # (768,) 벡터 하나는 크기 768인 1차원 벡터입니다. # print(embeddings[0].shape) # df의 text 컬럼에 있는 텍스트에 대응하는 인덱스를 저장합니다. df_to_index = df.set_index(["id"], drop=False) id_index = np.array(df_to_index.id.values) normalized_embeddings = embeddings.copy() # normalized_embeddings의 벡터를 L2 Norm(유클리드 거리 기준)으로 정규화합니다. # 이렇게 함으로써, 벡터의 길이를 1로 만들어 줍니다. 이것은 내적 기반의 유사도 검색에서 중요한 절차입니다. # 왜냐하면 정규화된 벡터들 사이의 내적은 두 벡터의 코사인 유사도와 동일하기 때문입니다. faiss.normalize_L2(normalized_embeddings) # faiss.IndexFlatIP는 faiss에서 제공하는 인덱싱 방법 중 하나로, 벡터들 간의 내적(inner product)을 기반으로 유사도 검색을 수행합니다. # 여기서 len(embeddings[0])는 임베딩 벡터의 차원수를 의미합니다. index_flat = faiss.IndexFlatIP(len(embeddings[0])) # faiss.IndexIDMap는 다른 인덱스 객체(여기서는 index_flat)와 함께 사용되어, # 벡터를 추가할 때 해당 벡터에 특정 ID를 매핑하여 저장하는 기능을 합니다. 이렇게 함으로써, 검색 시 벡터 대신 해당 벡터의 ID를 반환받을 수 있습니다. index = faiss.IndexIDMap(index_flat) # normalized_embeddings에 있는 벡터들을 index에 추가하면서, 각 벡터에 대한 ID를 id_index에서 가져와 함께 저장합니다. index.add_with_ids(normalized_embeddings, id_index) # 쿼리 함수를 정의합니다. def search(query, k):   # 쿼리도 앞에서 한것과 똑같이 인코딩하여 벡터로 변환하고 정규화 해야 합니다.   vector = model.encode([query])   faiss.normalize_L2(vector)   # k개 검색결과를 얻습니다. 인덱스와 유사도를 리턴받습니다.    similarities, ids = index.search(vector, k)     ids = ids[0]   similarities = similarities[0]   print(f'Query: {query}')   print()   results = df_to_index.loc[ids]  # id와 text 컬럼으로 구성된 데이터프레임에서 인덱스 ids를 위한 값을 가져옵니다.   results['similarity'] = similarities   # 해당 값과 쿼리 사이의 유사도입니다.     return results.reset_index(drop=True) # search 함수에 쿼리할 문자열과 리턴받을 검색 결과 개수를 적어줍니다. print(search('apple', 1)) print() print(search('pineapple', 3)) print() print(search('banana', 5)) print()

실행 결과입니다.

 

5개의 문장이 주어지며 각 문장에 id가 부여됩니다.

id                 text 0   0      I like a apple. 1   1     I like a orange. 2   2     I like a banana. 3   3  I like a pineapple. 4   4      I like a melon.

쿼리로 사용한 단어와 유사하다고 검색된 문장들의  유사도를 보여줍니다. 숫자가 높을수록 유사한 것입니다.

Query: apple    id             text  similarity 0   0  I like a apple.    0.714879 Query: pineapple    id                 text  similarity 0   3  I like a pineapple.    0.742127 1   0      I like a apple.    0.336482 2   2     I like a banana.    0.311522 Query: banana    id                 text  similarity 0   2     I like a banana.    0.695340 1   0      I like a apple.    0.380851 2   3  I like a pineapple.    0.306810 3   4      I like a melon.    0.244388 4   1     I like a orange.    0.234344