Vector 유사도 검색 - Elasticsearch

Elasticsearch를 이용해 vector 유사도 검색을 실험해 보자. 기존 역색인 검색과 비교해서 결과가 어떻게 달라지는지 확인해 보면 좋을 것 같다.

 

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

 

ES를 설치하고, 한글 형태소 분석기인 Nori를 설치하고(기존 역색인 검색과 비교하기 위해 설치) ES 데몬 인스턴스 만들고 접속하는 부분은 지난 글에 작성했기 때문에 넘어가겠다.

 

👇 이전 글 👇

구글 코랩에서 Elasticsearch 설치 및 실습

구글 코랩에서 Elasticsearch 설치 및 실습

 

1. 환경 설정

임베딩 생성을 위한 벡터 인코더를 설치하자.

!pip install sentence-transformers

그리고 데이터 전처리 부분이다.

일단 위키미디어로부터 데이터셋을 다운받아준다.

# 위키미디어로부터 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 이 모델을 사용하는데 한국어 임베딩 생성 가능한 어떤 모델을 사용해도 무방하다.

 

임베딩 생성하는 함수를 만들어주고, 다운로드한 데이터셋을 ES가 색인할 수 있는 파일 형태로 바꿔준다. 여기서 중간에 text, title 필드에 대해서 임배딩을 생성해 준다. 이렇게 생성된 임베딩을 따로 리스트로 만들어서 나중에는 실제로 이 값을 갖고 색인을 진행한다.

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'에 대한 임베딩을 계산하여 'embeddings' 필드에 추가
    json_data['embeddings_text'] = get_embedding(json_data['text']).tolist()
    json_data['embeddings_title'] = get_embedding(json_data['title']).tolist()

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

 

2. 색인 및 검색

이제 ES 색인 및 검색을 위한 공통 함수를 정의해 주자.

from elasticsearch import Elasticsearch, helpers
import json
import pprint as pp

def create_es_index(index, body):
    # 인덱스가 이미 존재하는지 확인
    if es.indices.exists(index=index):
        # 인덱스가 이미 존재하면 설정을 새로운 것으로 갱신하기 위해 삭제
        es.indices.delete(index=index)
    # 지정된 설정으로 새로운 인덱스 생성
    es.indices.create(index=index, body=body)

def delete_es_index(index):
    # 지정된 인덱스 삭제
    es.indices.delete(index=index)

def bulk_add(index, docs):
    # 대량 인덱싱 작업을 준비
    actions = [
        {
            '_index': index,
            '_source': doc
        }
        for doc in docs
    ]
    # Elasticsearch 헬퍼 함수를 사용하여 대량 인덱싱 수행
    return helpers.bulk(es, actions)

def sparse_retrieve(condition, index):
    # 지정된 인덱스에서 생성된 쿼리를 사용하여 검색 수행 (역색인을 이용한 일반 검색)
    return es.search(index=index, body=condition["query_body"], size=condition["size"], sort="_score")

def dense_retrieve(condition, index):
    # 벡터 유사도 검색에 사용할 쿼리 임베딩 가져오기
    query_embedding = get_embedding([condition["query"]])[0]

    # KNN을 사용한 벡터 유사성 검색을 위한 매개변수 설정
    knn = {
        "field": condition["field"],
        "query_vector": query_embedding.tolist(),
        "k": condition["size"],
        "num_candidates": 100
    }

    # 지정된 인덱스에서 벡터 유사도 검색 수행
    return es.search(index=index, knn=knn)

• sparse_retrieve: 이 함수는 기존 역색인 검색을 위한 함수다.
• dense_retrieve: 임배딩을 사용해서 검색하는 함수다. 이 함수는 쿼리가 들어왔을 때 쿼리에 대한 임배딩을 먼저 구하고, knn이라는 딕셔너리를 사용해서 검색을 위한 설정을 해준다.
  • field: 어떤 필드에서 검색할 것인가
  • query_vector: 실제 query에 대한 벡터는 무엇인가
  • k: 몇 개를 추출할 것인가
  • num_candidaties: ES에서는 내부적으로 shard를 사용하기 때문에 각 shard별로 100개의 candidates를 추출한 뒤에 그중에 한꺼번에 merge 해서 최종적으로 k개를 만든다.

이제 색인을 위한 mapping을 설정해 준다.
기존 코드와 유사하지만 아래 mappings 부분에 추가로 벡터 검색을 위해 embeddings_title, embeddings_text 항목이 추가되었다.

• type: 기존 역색인 때는 키워드라던지 text를 사용했다면 여기서는 dense_vector를 사용한다.
• dims: 몇 차원을 사용할 건지
• similarity: 유사도 계산을 어떤 함수로 할 것인지

settings > analysis > nori는 벡터 검색과 비교하기 위해 추가한 것이다.

# 색인을 위한 mapping 설정
setting = {
    "settings": {
        "analysis": {
            "analyzer": {
                "nori": {
                    "type": "custom",
                    "tokenizer": "nori_tokenizer",
                    "decompound_mode": "mixed",
                    "filter": ["nori_posfilter"]
                }
            },
            "filter": {
                "nori_posfilter": {
                    "type": "nori_part_of_speech",
                    # 어미, 조사, 구분자, 줄임표, 지정사, 보조 용언 등
                    "stoptags": ["E", "J", "SC", "SE", "SF", "VCN", "VCP", "VX"]
                }
            }
        }
    },
    "mappings": {
        "properties": {
            # 비교 테스트를 위해 meta field를 같이 색인
            "text": {"type": "text", "analyzer": "nori"},
            "title": {"type": "text", "analyzer": "nori"},
            "embeddings_title": {
                "type": "dense_vector",
                "dims": 768,
                "index": True,
                "similarity": "l2_norm"
            },
            "embeddings_text": {
                "type": "dense_vector",
                "dims": 768,
                "index": True,
                "similarity": "l2_norm"
            }
        }
    }
}

위에서 세팅한 설정으로 인덱스와 데이터를 추가해 준다.

# 'setting'으로 설정된 내용으로 'test' 인덱스 생성
create_es_index("test", setting)

# 'test' 인덱스에 대량 색인화 수행
# 각 문서는 'embbedings' 라는 필드를 가짐
ret = bulk_add("test", index_docs)

# 결과 출력
print(ret)

역색인을 사용하는 검색 예제

query로 '문재인의 친구'를 주고 text 필드에서 검색하는 코드다.

condition_retrieve = {
    "query_body": {
        "query": {
            "match": {
                "text": {
                    "query": "문재인의 친구"
                }
            }
        }
    },
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = sparse_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

기본적으로 ES는 OR operator를 사용하기 때문에 '문재인의 친구'가 전부 나오기보다는 각 단어가 나와서 추출된 경우가 많다.

그래서 operator를 AND로 수정해 주고 실행해 보자

# 역색인을 사용하는 검색 예제
condition_retrieve = {
    "query_body": {
        "query": {
            "match": {
                "text": {
                    "query": "문재인의 친구",
                    "operator": "AND"
                }
            }
        }
    },
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = sparse_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

기존대비 검색 결과 양이 줄어든 걸 확인할 수 있다.

Vector 유사도 사용한 검색 예제

이제 vector 유사도를 사용해서 검색해 보자.

embeddings_text 필드에서 찾아보자.

condition_retrieve = {
    "field": "embeddings_text",
    "query": "문재인의 친구",
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = dense_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

이 결과는 만약에 '문재인의 친구'와 유사한 의미를 가진 단어가 있다면 상위로 추출된 검색 결과들이다.

역색인을 사용하는 검색 예제 - title field 사용

지금까지 text 필드에 대해서만 진행했던걸 title 필드에서 검색해 보자.
사실 기존에 text 필드는 보면 사이즈가 너무 길기 때문에 임베딩을 생성했을 때 성능이 그렇게 높게 나오기 힘들다. 그래서 비교적 짧은 title로 시도해 보자.

condition_retrieve = {
    "query_body": {
        "query": {
            "match": {
                "title": {
                    "query": "문재인의 친구"
                }
            }
        }
    },
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = sparse_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

현재 query로 '문재인의 친구'를 검색했을 때 검색 결과가 하나도 나오지 않았다. 이 원인은 실제 title 애 '문재인', '친구'를 갖는 title이 없기 때문이다.

Vector 유사도 사용한 검색 예제 - title field 사용

이제 벡터 유사도로 검색해 보자.

condition_retrieve = {
    "query": "문재인의 친구",
    "field": "embeddings_title",
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = dense_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

실제로 '노무현'이라는 title이 최상으로 나왔는데 이는 '문재인', '친구'라는 단어는 안 나왔지만 의미상으로 뭔가 임베딩이 유사하게 벡터공간에 있기 때문에 해당 결과를 가져왔다고 생각한다.

 

이제 조금 더 복잡한 쿼리를 해보자.

역색인 검색

title에서 기존 역색인 방법으로 다음과 같이 검색하면 기본적으로 OR 검색이기 때문에 '우리나라', '열', '여섯', '대통령' 이런 단어들이 포함된 title을 추출할 것이다.

condition_retrieve = {
    "query_body": {
        "query": {
            "match": {
                "title": {
                    "query": "우리나라 열 여섯번째 대통령이 누구야?"
                }
            }
        }
    },
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = sparse_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

검색 결과를 보면 '나라', '대통령' 이런 단어가 들어간 title을 뽑아온 걸 볼 수 있다.
단어를 보고 하나라도 추출된 게 있다면 뽑아오는걸 확인했다. 아마 AND operator로 실행했으면 검색 결과가 안 나오지 않았을까 싶다.

Vector 유사도

Vector 유사도로 검색해 보자.

condition_retrieve = {
    "query": "우리나라 열 여섯번째 대통령이 누구야?",
    "field": "embeddings_title",
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = dense_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

이렇게 나오면 '대한민국 제16대 대통령 선거'가 최상위로 뽑힌 걸 볼 수 있다. 역색인 방법보다도 더 의미적으로 더 가까운 결과를 내는걸 볼수있다.

역색인을 사용할 때 오탈자

오탈자 같은 경우에는 어떻게 될까? 대통령을 대통려라고 오타를 내보자.

condition_retrieve = {
    "query_body": {
        "query": {
            "match": {
                "title": {
                    "query": "우리나라 열 여섯번째 대통려이 누구야?"
                }
            }
        }
    },
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = sparse_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

 

실제로 제목에 '대통령'이 들어간 title이 있을 테지만 지금은 나오지 않고 있다.

Vector 유사도 사용할 때 오탈자

Vector 유사도를 사용할 때 오탈자도 확인해 보자.

condition_retrieve = {
    "query": "우리나라 16대 대통려이 누구야?",
    "field": "embeddings_title",
    "size": 5  # Specify the number of documents to retrieve
}
search_result_retrieve = dense_retrieve(condition_retrieve, "test")

# 결과 출력
for rst in search_result_retrieve['hits']['hits']:
    print('score:', rst['_score'], 'source::', rst['_source'])

검색 결과

아까와 유사하게 잘 나오는 걸 확인할 수 있다.