Vector 유사도 검색 - Faiss

이번 실습에서는 Faiss를 이용하여 vector 유사도 검색을 실험해 보자.

이번 글에서는 wikimedia kowiki로 wikimedia에서 제공하는 한국어 데이터셋을 이용할 거다.

 

1. 환경 설정

먼저 Faiss 패키지를 설치해 준다. gpu와 cpu를 사용하는 패키지가 서로 다른데 나는 구글 코랩에서 gpu를 사용할 거기 때문에 gpu 버전으로 다운받아줬다.

# GPU를 사용하는 버전의 Faiss Python 패키지 설치
!pip install faiss-gpu

그리고 임베딩 생성을 위한 벡터 인코더로 sentence-transformers를 설치해 준다.

!pip install sentence-transformers

2. 데이터 전처리

이제 데이터를 다운받고 전처리를 해주자.

# 위키미디어로부터 kowiki 데이터를 다운로드 받음
!wget https://dumps.wikimedia.org/kowiki/latest/kowiki-latest-pages-articles1.xml-p1p82407.bz2

# 위키데이터의 노이즈를 제거하고 json 형태로 반환하는 코드를 참조
!git clone https://github.com/attardi/wikiextractor.git

# 다운로드 받은 샘플 위키 데이터를 전처리하여 검색의 입력으로 사용
# 결과는 elastic 폴더에 'extract_result/AA,AB,AC.../wiki_00..99'라는 새로운 폴더에 저장된다.(용량이 비슷하게 나눠서 저장됨)
# 변환결과 wiki_00 파일의 내용 샘플  {"id": "5", "revid": "641228", "url": "https://ko.wikipedia.org/wiki?curid=5", "title": "\uc9c0\...\ud130", "text": "\uc81c\...\ub2e4."}
!python -m wikiextractor.wikiextractor.WikiExtractor kowiki-latest-pages-articles1.xml-p1p82407.bz2 --json -o extract_result

다운을 다 받았으면 전처리를 해주자. 데이터 전처리는 임베딩을 만들어주는 거라고 생각하면 된다.

 

우선 SentenceTransformer 모델을 초기화해 준다. 지금은 hunkim/sentence-transformer-klue 이 모델을 사용하는데 한국어 임베딩 생성 가능한 어떤 모델을 사용해도 무방하다.

import json
from sentence_transformers import SentenceTransformer

# Sentence Transformer 모델 초기화
model = SentenceTransformer("hunkim/sentence-transformer-klue")

def get_embedding(sentences):
    # 입력 문장을 인코딩하여 임베딩을 얻음
    return model.encode(sentences)

wiki_dump_json_file = '/content/extract_result/AA/wiki_00'
# 'wiki_dump_json_file'에 있는 JSON 파일 읽어들여 index_docs에 저장
index_docs = []

for line in open(wiki_dump_json_file, encoding="utf-8"):
    # JSON 데이터를 읽어들여 파이썬 딕셔너리로 변환
    json_data = json.loads(line)

    # text와 title에 대한 임베딩을 계산하여 해당 필드에 추가
    json_data['embeddings_title'] = get_embedding(json_data['title']).tolist()
    json_data['embeddings_text'] = get_embedding(json_data['text']).tolist()

    # 색인할 문서 목록에 추가
    index_docs.append(json_data)

Faiss는 특정 임베딩 필드 하나만 저장할 수 있기 때문에 아래와 같이 임베딩 필드만 추출해서 변수에 저장해 주도록 하자.

title_docs = [item["embeddings_title"] for item in index_docs]
text_docs = [item["embeddings_text"] for item in index_docs]

3. 색인 및 검색

이제 실제로 Faiss를 이용해서 색인과 검색을 해보자.

 

우선 테스트로 가장 간단한 Index type인 FlatL2를 사용해 보도록 하자. Faiss는 다양한 ANN 알고리즘을 제공하는데 Flat, IVF, HNSW, PQ 등 다양하게 제공한다. 여기서는 우선 문서가 몇 개 안 되니까 기본인 FlatL2를 사용해 본 것이다.

 

아까 저장해 놨던 변수 두 개를 사용해서 색인을 진행해 보자.

하나는 title field, 하나는 text field로 지정해 줬다.

import faiss
import numpy as np

index_title = faiss.IndexFlatL2(768)
index_title.add(np.array(title_docs).astype('float32'))

index_text = faiss.IndexFlatL2(768)
index_text.add(np.array(text_docs).astype('float32'))

그럼 이제 실제로 쿼리를 날려보자.
실제 쿼리가 있으면 임베딩을 생성하고 그 임베딩을 search 인자로 넣어주면 된다. 두 번째 인자 5는 다섯 개를 추출한다는 의미다.

# 결과로 offset과 score가 나오는데 원본 문서와의 매칭은 offset 값을 이용

query_str = "문재인의 친구"
query_emb = get_embedding([query_str])
scores, offsets = index_title.search(query_emb, 5)

for i,offset in enumerate(offsets[0]):
    print(f'{offset}, {scores[0][i]} {index_docs[offset]}')

검색 결과

 

결과가 다음과 같이 나왔다.

 

제목에서 검색을 했는데 title에 '문재인', '친구'가 안 나왔지만 '노무현'이 최상으로 출력된 걸 볼 수 있다.

그런데 Faiss는 문서 전체를 색인하는 게 아니라 임베딩만 색인하기 때문에 이 return값인 offset을 이용해서 기존 문서와 매칭해줘야 한다.그래서 for loop를 돌 때도 offset을 이용해서 실제 문서 index_docs 변수에서 offset으로 접근해서 원래 field
에 있던 내용을 가져오게 된다.

 

💡 Faiss에서 offset은 색인(index)에 저장된 데이터의 위치를 나타내는 값이다. Faiss는 원본 데이터를 직접 색인하지 않고, 데이터(예: 문서, 문장)의 임베딩(벡터)을 생성해 이 벡터를 색인에 저장한다. 이 과정에서 Faiss는 각 벡터의 색인 번호(또는 위치)를 기록하고, 이를 통해 결과를 원본 데이터와 매칭할 수 있다.

 

다른 질의로 질문해 보자.

query_str = "대한민국 16대 대통령이 누구야?"
query_emb = get_embedding([query_str])
scores, offsets = index_title.search(query_emb, 5)

for i,offset in enumerate(offsets[0]):
    print(f'{offset}, {scores[0][i]} {index_docs[offset]}')

검색 결과

이때도 잘 찾아오는 걸 확인할 수 있다.

text 필드에서도 검색해 보자.

query_str = "문재인의 친구"
query_emb = get_embedding([query_str])
scores, offsets = index_text.search(query_emb, 5)

for i,offset in enumerate(offsets[0]):
    print(f'{offset}, {scores[0][i]} {index_docs[offset]}')

검색 결과

 

사실 title은 사용하지 않지만 offset 매칭 과정에서 나온 거다. 아마 title을 '노무현'으로 한 데이터의 text 내용 중에서 '문재인의 친구'와 유사한 의미를 갖는 문구들이 있을 거라고 예상해 본다.

다른 질의로 한 번 더 실행해 보자.

query_str = "대한민국 16대 대통령이 누구야?"
query_emb = get_embedding([query_str])
scores, offsets = index_text.search(query_emb, 5)

for i,offset in enumerate(offsets[0]):
    print(f'{offset}, {scores[0][i]} {index_docs[offset]}')

검색 결과

 

역시나 잘 찾아온 걸 볼 수 있다.

 

이번에도 오탈자를 내보자. 대통령을 '대통려'로 오타를 내보고 질의해 봤다.

query_str = "대한민국 16대 대통려이 누구야?"
query_emb = get_embedding([query_str])
scores, offsets = index_text.search(query_emb, 5)

for i,offset in enumerate(offsets[0]):
    print(f'{offset}, {scores[0][i]} {index_docs[offset]}')

검색 결과

 

결과를 보면 잘 찾아오는 걸 볼 수 있다.