[Python] Faiss 활용한 Vector 유사도 측정 및 검색

faiss : 페이스북에서 C++로 제작한 Vector 자료 구조를 색인화하고 유사성을 검색하는 모듈이다.

CPU 기반의 모듈(pip install faiss-cpu )과 GPU (pip install faiss-gpu ) 모듈이 각각이 있어 환경에 맞는 모듈을 설치한다.

 

# pip install faiss-cpu
import faiss

def create_index_kmean(vector, size, saved_file=None):
    dimension = 768
    index_obj = faiss.IndexIDMap( faiss.IndexFlatIP(dimension) )
    xidx = np.array( range(0, size), dtype='int64' )
    index_obj.add_with_ids(vector, xidx)
    
    if savefile:
    	faiss.write_index(index_obj, savefile)
    
    return index_obj

def create_index_cosine(vector, size, saved_file=None):
    dimension = 768
    index_obj = faiss.IndexIDMap( faiss.IndexFlatIP(dimension) )
    xidx = np.array( range(0, size), dtype='int64')
    faiss.normalize_L2(vector) 	 # cosine 유사도에서는 normalize_L2 처리.
    index_obj.add_with_ids(vector, xidx)
    
    if savefile:
    	faiss.write_index(index_obj, savefile)
    
    return index_obj
    
#인덱스 생성
index_obj = create_index_kmean( database_vector, size=len(database_vector) )

k = 5 # 최대 5개의 서치 결과를 반환하도록 설정.
distance, index = index_obj.search(query_vec, k) 

print(distance[0:])
print(index[0:])