[LLM] 임베딩 모델로 데이터 의미 압축하기 (3) - 의미, 키워드, 하이브리드 검색

 

 

사용자의 쿼리를 바탕으로 이에 걸맞은 정보를 얻기 위해서는 어떤 방식의 검색 알고리즘을 사용하는 것이 좋을까? 이번 시간에는 임베딩 변환 후 유사도 기반의 의미 검색과 단순히 키워드로 같은 문자열을 찾아내는 키워드 검색 중 하나인 BM25, 그리고 이들을 합친 하이브리드 검색에 대해 구현해 보자.

 

---

의미 검색

 

사실 수식으로 따지면 임베딩이 차지하는 비율이 높기 때문에 굉장히 복잡한 구현이 필요하지만, 우리는 임베딩 라이브러리인 sentence-transformer와 벡터 연산 라이브러리인 faiss를 통해 간단히 구현할 수 있다.

 

from datasets import load_dataset
from sentence_transformers import SentenceTransformer

klue_mrc_dataset = load_dataset('klue', 'mrc', split='train') # 데이터셋 불러오기
sentence_model = SentenceTransformer('snunlp/KR-SBERT-V40K-klueNLI-augSTS') # 모델 불러오기

 

이어서 데이터 중 1000개 만을 가져와 선언한 SBERT 모델을 이용하여 바로 임베딩을 진행한 후 차원을 살펴본다.

 

klue_mrc_dataset = klue_mrc_dataset.train_test_split(train_size=1000, shuffle=False)['train']
embeddings = sentence_model.encode(klue_mrc_dataset['context'])
embeddings.shape
# 출력 결과
# (1000, 768)

 

인덱스를 선언 후에 방금 생성한 임베딩 벡터들을 인덱스에 저장한다.

 

import faiss
# 인덱스 만들기
index = faiss.IndexFlatL2(embeddings.shape[1])
# 인덱스에 임베딩 저장하기
index.add(embeddings)

 

이제 쿼리를 생성해 보자. 문장 쿼리 또한 임베딩 변환을 거친 후에 이와 유사한 벡터를 찾아 상위 3개를 뽑아 보았다.

 

query = "이번 연도에는 언제 비가 많이 올까?"
query_embedding = sentence_model.encode([query])
distances, indices = index.search(query_embedding, 3)

for idx in indices[0]:
  print(klue_mrc_dataset['context'][idx][:50])

# 출력 결과
# 올여름 장마가 17일 제주도에서 시작됐다. 서울 등 중부지방은 예년보다 사나흘 정도 늦은   (정답)
# 연구 결과에 따르면, 오리너구리의 눈은 대부분의 포유류보다는 어류인 칠성장어나 먹장어, 그 (오답)
# 연구 결과에 따르면, 오리너구리의 눈은 대부분의 포유류보다는 어류인 칠성장어나 먹장어, 그 (오답)

 

출력이 잘 된 것을 확인할 수 있었다. 하지만 어디까지나 의미 검색이라 함은 임베딩 벡터의 유사도를 비교하는 것이다. 따라서 우리가 예상치 못하는 다음과 같은 결과도 발생할 수 있다.

 

query = klue_mrc_dataset[3]['question'] # 로버트 헨리 딕이 1946년에 매사추세츠 연구소에서 개발한 것은 무엇인가?
query_embedding = sentence_model.encode([query])
distances, indices = index.search(query_embedding, 3)

for idx in indices[0]:
  print(klue_mrc_dataset['context'][idx][:50])

# 출력 결과
# 태평양 전쟁 중 뉴기니 방면에서 진공 작전을 실시해 온 더글러스 맥아더 장군을 사령관으로 (오답)
# 태평양 전쟁 중 뉴기니 방면에서 진공 작전을 실시해 온 더글러스 맥아더 장군을 사령관으로 (오답)
# 미국 세인트루이스에서 태어났고, 프린스턴 대학교에서 학사 학위를 마치고 1939년에 로체스 (정답)

 

키워드 검색

 

키워드 검색은 의미 검색과 달리 동일한 키워드가 많이 포함될수록 유사도를 높게 평가하는 검색 방식을 말한다. 대표적인 키워드 검색 방식으로는 TF-IDT의 단점을 보완한 BM25 가 쓰인다. 이는 대표적인 검색 엔진인 Elasticsearch의 기본 알고리즘으로 사용된다.

 

BM25 구현

 

BM25의 전체 수식은 위와 같다.

 

$ IDF(q_i) $

특정 단어나 토큰 $ q_i $ 가 전체 문서에서 얼마나 자주 등장했는지를 나타낸다. $ n(q_i) $ 가 작다면 IDF 값이 감소하여 해당 문서의 중요도가 떨어지고, 반대로 $ n(q_i) $가 크다면 문서의 중요도가 증가한다.

 

$ f(q_i, D) $

특정 문서 $ D $ 에서 $ q_i $ 가 등장하는 횟수를 의미한다. 이 값은 아무리 커져도 분모와 분자 둘 다에 있기 때문에 계속해서 증가하지 않고 $ k_1 + 1 $ 이라는 포화 상태에 근사하게 함으로써 특정 문서의 중요도가 비정상적으로 증가하는 것을 막는다.

 

$ |D|, avgdl $

차례대로 특정 문서의 길이와 전체 문서의 길이의 평균을 뜻한다. 이 값의 대소 관계 또한 중요도에 영향을 미친다.

 

위 수식을 토대로 BM25를 구현하면 아래와 같다.

 

import math
import numpy as np
from typing import List
from transformers import PreTrainedTokenizer
from collections import defaultdict

class BM25:
  def __init__(self, corpus:List[List[str]], tokenizer:PreTrainedTokenizer):
    self.tokenizer = tokenizer
    self.corpus = corpus
    self.tokenized_corpus = self.tokenizer(corpus, add_special_tokens=False)['input_ids']
    self.n_docs = len(self.tokenized_corpus)
    self.avg_doc_lens = sum(len(lst) for lst in self.tokenized_corpus) / len(self.tokenized_corpus)
    self.idf = self._calculate_idf()
    self.term_freqs = self._calculate_term_freqs()

# 각 토큰이 몇개의 문서에 등장하는지 집계
  def _calculate_idf(self):
    idf = defaultdict(float)

	# 토큰화된 문서들의 집합에서
    for doc in self.tokenized_corpus:
    	# 고유한 토큰들만 추출하여 해당 토큰을 가진 문서가 몇개인지 저장
      	for token_id in set(doc):
        idf[token_id] += 1
    # 그 정보를 바탕으로 idf 계산
    for token_id, doc_frequency in idf.items():
    	idf[token_id] = math.log(((self.n_docs - doc_frequency + 0.5) / (doc_frequency + 0.5)) + 1)
    return idf

# 특정 문서에서 특정 토큰의 등장 횟수 저장
  def _calculate_term_freqs(self):
    term_freqs = [defaultdict(int) for _ in range(self.n_docs)]
    for i, doc in enumerate(self.tokenized_corpus):
      for token_id in doc:
        term_freqs[i][token_id] += 1
    return term_freqs
	
# 위 두 메서드를 바탕으로 BM25값 계산
  def get_scores(self, query:str, k1:float = 1.2, b:float=0.75):
    query = self.tokenizer([query], add_special_tokens=False)['input_ids'][0]
    scores = np.zeros(self.n_docs)
    for q in query:
      idf = self.idf[q]
      for i, term_freq in enumerate(self.term_freqs):
        q_frequency = term_freq[q]
        doc_len = len(self.tokenized_corpus[i])
        score_q = idf * (q_frequency * (k1 + 1)) / ((q_frequency) + k1 * (1 - b + b * (doc_len / self.avg_doc_lens)))
        scores[i] += score_q
    return scores
	
# 상위 점수 k개 반환하는 코드
  def get_top_k(self, query:str, k:int):
    scores = self.get_scores(query)
    top_k_indices = np.argsort(scores)[-k:][::-1]
    top_k_scores = scores[top_k_indices]
    return top_k_scores, top_k_indices

 

이어서 구현한 BM25의 성능을 테스트해 보자. 앞서 의미 검색의 장단점을 나타냈던 예시들을 BM25를 활용하여 똑같이 실행시켜 보면, 예상대로 이번에는 오답이 나오는 것을 발견할 수 있다.

 

# BM25 검색 준비
bm25 = BM25(klue_mrc_dataset['context'], tokenizer)

query = "이번 연도에는 언제 비가 많이 올까?"
_, bm25_search_ranking = bm25.get_top_k(query, 100)

for idx in bm25_search_ranking[:3]:
  print(klue_mrc_dataset['context'][idx][:50])

# 출력 결과
# 갤럭시S5 언제 발매한다는 건지언제는 “27일 판매한다”고 했다가 “이르면 26일 판매한다 (오답)
# 인구 비율당 노벨상을 세계에서 가장 많이 받은 나라, 과학 논문을 가장 많이 쓰고 의료 특 (오답)
# 올여름 장마가 17일 제주도에서 시작됐다. 서울 등 중부지방은 예년보다 사나흘 정도 늦은  (정답)

 

하지만 의미 검색에서 오답이 나왔던 예시를 돌려 보면 아래와 같이 정답이 나오는 것을 볼 수 있다!

 

query = klue_mrc_dataset[3]['question']  # 로버트 헨리 딕이 1946년에 매사추세츠 연구소에서 개발한 것은 무엇인가?
_, bm25_search_ranking = bm25.get_top_k(query, 100)

for idx in bm25_search_ranking[:3]:
  print(klue_mrc_dataset['context'][idx][:50])

# 출력 결과
# 미국 세인트루이스에서 태어났고, 프린스턴 대학교에서 학사 학위를 마치고 1939년에 로체스 (정답)
# ;메카동(メカドン)                                                      (오답)
# :성우 : 나라하시 미키(ならはしみき)
# 길가에 버려져 있던 낡은 느티나
# ;메카동(メカドン)                                                      (오답)
# :성우 : 나라하시 미키(ならはしみき)
# 길가에 버려져 있던 낡은 느티나

 

상호 순위 조합

 

그렇다면 의미 검색과 키워드 검색을 상황에 맞게 적절하게 사용하는 방법은 어떨까? 그렇게 해서 등장한 방법이 상호 순위 조합 개념이다. 두 집단에서의 순위를 동시에 반영하여 최종 순위를 집계한다.

 

from collections import defaultdict

def reciprocal_rank_fusion(rankings:List[List[int]], k=5):
    rrf = defaultdict(float)
    for ranking in rankings:
        for i, doc_id in enumerate(ranking, 1):
            rrf[doc_id] += 1.0 / (k + i)
    return sorted(rrf.items(), key=lambda x: x[1], reverse=True)
    
rankings = [[1, 4, 3, 5, 6], [2, 1, 3, 6, 4]]
reciprocal_rank_fusion(rankings)

# [(1, 0.30952380952380953),
#  (3, 0.25),
#  (4, 0.24285714285714285),
#  (6, 0.2111111111111111),
#  (2, 0.16666666666666666),
#  (5, 0.1111111111111111)]

 

하이브리드 검색

 

자 이제 모두 준비가 되었다. 한 쿼리에 대해 의미 검색, 키보드 검색을 진행 후 상호 순위 조합 메서드를 사용하여 결과를 내보내 보자.

 

# 의미 검색 반환
def dense_vector_search(query:str, k:int):
  query_embedding = sentence_model.encode([query])
  distances, indices = index.search(query_embedding, k)
  return distances[0], indices[0]

# 의미 검색 랭킹과 BM25 랭킹 반환
def hybrid_search(query, k=20):
  _, dense_search_ranking = dense_vector_search(query, 100)
  _, bm25_search_ranking = bm25.get_top_k(query, 100)
	
	# 상호 순위 조합 함수로 결과 출력
  results = reciprocal_rank_fusion([dense_search_ranking, bm25_search_ranking], k=k)
  return results
  
  query = "이번 연도에는 언제 비가 많이 올까?"
print("검색 쿼리 문장: ", query)
results = hybrid_search(query)
for idx, score in results[:3]:
  print(klue_mrc_dataset['context'][idx][:50])

print("=" * 80)
query = klue_mrc_dataset[3]['question'] # 로버트 헨리 딕이 1946년에 매사추세츠 연구소에서 개발한 것은 무엇인가?
print("검색 쿼리 문장: ", query)

results = hybrid_search(query)
for idx, score in results[:3]:
  print(klue_mrc_dataset['context'][idx][:50])

# 출력 결과
# 검색 쿼리 문장:  이번 연도에는 언제 비가 많이 올까?
# 올여름 장마가 17일 제주도에서 시작됐다. 서울 등 중부지방은 예년보다 사나흘 정도 늦은  (정답)
# 갤럭시S5 언제 발매한다는 건지언제는 “27일 판매한다”고 했다가 “이르면 26일 판매한다  (오답)
# 연구 결과에 따르면, 오리너구리의 눈은 대부분의 포유류보다는 어류인 칠성장어나 먹장어, 그 (오답)
# ================================================================================
# 검색 쿼리 문장:  로버트 헨리 딕이 1946년에 매사추세츠 연구소에서 개발한 것은 무엇인가?
# 미국 세인트루이스에서 태어났고, 프린스턴 대학교에서 학사 학위를 마치고 1939년에 로체스 (정답)
# 1950년대 말 매사추세츠 공과대학교의 동아리 테크모델철도클럽에서 ‘해커’라는 용어가 처음 (오답)
# 1950년대 말 매사추세츠 공과대학교의 동아리 테크모델철도클럽에서 ‘해커’라는 용어가 처음 (오답)

 

두 경우 모두에서 원했던 정답이 나왔음을 확인할 수 있었다.