파이썬 구현으로 알아보는 TF-IDF와 BM25

개인적으로 유사한 컨텐츠를 찾는 검색을 할 때 1) LM embedding 기반의 semantic semantic search와 2)기존 키워드 기반 둘 다 중요하다고 생각한다. 

 

전통적인 머신러닝은 비교적 간단하지만 그만큼 강력하다고 생각한다. 오늘은 간단하게 파이썬으로 구현해보면서 TF-IDF와 BM25에 대하여 알아보자. 

 

모든 코드는 google colab에서 동일하게 실행시킬 수 있다.

 

---

0. data preprocess 

데이터셋은 huggingface에서 문서 데이터를 불러왔다. 한국 판결에 대한 데이터셋을 1000개만 불러와보자. 

import datasets
sample_dataset = datasets.load_dataset("joonhok-exo-ai/korean_law_open_data_precedents", split="train[:10%]")

sample_dataset = sample_dataset['판결요지']
print(sample_dataset[0]) # 회사소유의 부동산을 회사대표자인 개인이 계약당사자로서 매도하고 다시 회사대표자 자격으로써 한 소유권이전등기는 원인없는 등기이다.

my_datasets = pd.Series(sample_dataset[:1000]).apply(lambda x: preprocess(x))

 

 

문서를 토큰처럼 쪼개는 전처리(preprocess)과정을 거쳐야 한다. 이 과정에서 조사나 특수문자를 없앨수도 있다. 동사를 명사화할 수도 있다. 코드 예시에서는 konlpy 형태소 분석기의 Okt를 사용했다.(원래는 Mecab을 사용하는 편이다.) colab에서 konlpy를 설치할 수 있다. 

 

! pip3 install konlpy

 

일련의 preprocess가 필요한 이유는 키워드 기반으로 문서를 벡터화할거기 때문이다. 그러기 위해선 키워드가 잘 정제되어야 한다.

from konlpy.tag import Mecab, Okt
from typing import List

okt = Okt()
def preprocess(row: str) -> List:
  return okt.morphs(row)
  
  
preprocess(sample_dataset[0])[:10] # -> ['이전', '시간', '에',]...

 

 

1000개의 문서만 전처리해서 사용해보자. 

import pandas as pd 

my_datasets = pd.Series(sample_dataset[:1000]).apply(lambda x: preprocess(x))

1.  TF-IDF 

확보해둔 문서들을 키워드 기반으로 벡터화해보자. TF-IDF는 문서 내 단어의 분포를 기반으로 문서를 벡터화할 수 있다. 

먼저 TF다. TF는 각 문서들에서 키워드가 몇번 나타났는지를 단순히 count한다. 

from collections import defaultdict
from typing import Dict 

def get_TF(doc: list) -> Dict:
  word_dict = defaultdict(int)
  for word in doc: 
    word_dict[word] += 1 #ex) {'안녕': 1, ...}

  return word_dict

 

문서들의 TF (Term Frequency)를 구하면 다음과 같다. 

tf_datasets = my_datasets.apply(lambda x: get_TF(x))
# 0번째 문서의 tf를 확인해보자.
tf_datasets[0] # defaultdict(int,{'이전': 1, '시간': 1, '에': 10, '우리': 5, '가': 5, ..}

 

 

다음은 DF(Document Frequency)다. TF가 문서를 기반으로 term의 딕셔너리를 형성했다면 이번엔 '단어'를 기반으로 형성한다. 해당 단어가 지금까지 몇개의 문서(document)에서 등장했는지(frequency)를 센다. 

def get_DF(docs_tf: dict) -> Dict:
  word_dict = defaultdict(int)
  total_words = []
  for doc_tf in docs_tf: 
    total_words.extend(doc_tf.keys()) # 1000개 문서의 모든 word를 불러온다. 
  
  for word in total_words:
    for doc_tf in docs_tf:
      if word in doc_tf.keys():
        word_dict[word] += 1 #word가 문서에 존재한다면 count +=1

  return word_dict

 

우리의 1000개 문서에 대해 DF를 계산해보자. 

df_datasets = get_DF(tf_datasets)

 

 

IDF는 DF의 역수다. DF가 0일 경우 무한대로 발산하므로 분모에 1을 더해준다. 

idf 수식

 

구현에서는 추가로 log 내에 +1을 추가했다. idf값이 음수로 나오는 것을 방지하기 위해서다. 

import numpy as np 


def get_IDF(docs_df: dict, n_docs : int):
  prevent_minus = 1 # idf값이 음수가 나와서 양수로 바꿔주기 위해 추가함.
  return {word: np.log(n_docs / (count + 1) + prevent_minus) for word, count in docs_df.items()}
  
idf_datasets = get_IDF(df_datasets, len(my_datasets))

 

이제 my_datasets의 문서 1000개를 TF-IDF로 표현할 수 있게 되었다. TF-IDF는 tf * idf로 표현할 수 있다. 

def calc_TF_IDF(doc_tf: dict, docs_idf: dict):
  doc_tf_idf = defaultdict(dict)
  doc_tf_idf = {k : 0 for k in docs_idf.keys()}

  for word, tf in doc_tf.items():
    doc_tf_idf[word] = tf * docs_idf[word]
  
  return doc_tf_idf

tf_idf_datasets = tf_datasets.apply(lambda x: calc_TF_IDF(x, idf_datasets))

 

문서별 tf_idf를 다 구했다! 각 문서들을 TF-IDF 기반으로 벡터화했다. 

final_tf_idf = tf_idf_datasets.apply(pd.Series).fillna(0)
final_tf_idf.head()

 

결과물은 다음과 같다. row는 총 1000개, column은 총 5,011개다. 즉, 문서를 TF-IDF로 표현한 결과는 1000개의 문서를 5,011개의 단어 빈도로 표현한 것과 동일하다. 

 

 

2. BM25

BM25 스코어를 구하는 공식은 다음과 같다.

• f(t,d): 위에서 구한 TF 
• ∣d∣: 문서 길이 (단어 수)
• avg: 전체 문서의 평균 길이
• k,b: 하이퍼파라미터

 

 

위에서 구한 TF와 IDF를 이용해서 BM25를 구할 수 있다. 하이퍼파라미터인 k= 1.5, b = 0.75로 가정해서 수식을 파이썬으로 옮겨보자.

def get_bm25_score(my_query: list, target_doc_idx: int, avg_doc_len: float, k:float = 1.5, b:float = 0.75):
  cur_tf = tf_datasets[target_doc_idx] 
  doc_len = len(cur_tf)
  score = 0

  for word in my_query:
    numerator = cur_tf[word] * (k+1)
    denominator = cur_tf[word] + k * (1 - b + b * doc_len / avg_doc_len)
    score += idf_datasets[word] * numerator / denominator

  return score

 

샘플로 문장 하나와 1번째 문서간의 BM25 score를 구해보자. 

avg_doc_len = np.mean([len(tf) for tf in tf_datasets])
get_bm25_score(['부동산', '을', '개인'], target_doc_idx = 1, avg_doc_len = avg_doc_len)
# 3.66

 

 

이제 주어진 쿼리와 우리의 1000개 문서간 score를 구해서, 가장 유사한 문서를 return해보자. 

def get_bm25_rank(my_query: list):
  bm25_rank = []
  assert len(my_datasets) == 1000
  for i in range(len(my_datasets)): 
    score = get_bm25_score(my_query, target_doc_idx = i, avg_doc_len = avg_doc_len)
    bm25_rank.append((i, score)) #i번쨰 문서, i번째 문서와 현재 쿼리간의 score를 담음. 
  
  bm25_rank.sort(key = lambda x: x[1], reverse = True)
  return bm25_rank

 

 

0번째 문서와 가장 유사한 문서는 무엇일까? 당연하지만 가장 높은 rank는 0번째 문서다.

print(my_datasets[0]) #['회사', '소유', '의', '부동산', '을', '회사', '대표자', '인', '개인', '이', '계약', '당사자', '로서', '매도', '하고', '다시', '회사', '대표자', '자격', '으로써', '한', '소유권', '이전', '등기', '는', '원인', '없는', '등기', '이다', '.']
rank = get_bm25_rank(my_datasets[0]) 

print(rank[:5]) # 유사도 상위 10개 
# [(0, np.float64(34.028452630962626)),
# (717, np.float64(23.099057481033263)),
# (145, np.float64(13.042570627170605)),
# (801, np.float64(10.621484865360282)),
# (668, np.float64(10.016698462103971)),

 

두번째 rnk가 높았던 717번째 문서를 보자. 비슷한가?

print(my_datasets[0])
#['회사', '소유', '의', '부동산', '을', '회사', '대표자', '인', '개인 ...

print(my_datasets[717])
# ['회사', '이름', '과', '대표이사', '의', '이름', '만', '기', '명', '날인', '되고' ...]

 

 

 

간단하게 쿼리와 문서간의 랭킹을 구했다. 다만 TF-IDF기반이기 때문에, 기존 1000개의 문서에 없던 모르는 단어가 쿼리에 있을 경우는 제대로 랭킹을 매길 수 없을 것이다. 

 

 

 

---

참고로 TF IDF는 spark에 이미 구현되어있다.