[Sklearn] 파이썬 TF-IDF 구하기, 코사인 유사도로 비슷한 문서 찾기

파이썬 사이킷런 유사 문서 탐색 예제

안녕하세요. 이번 글에서는 파이썬 scikit-learn 패키지를 이용하여

문서의 TF-IDF 벡터를 구해보고, 구한 벡터 간의 코사인 유사도를 기반으로

비슷한 문서를 찾는 과정에 대하여 다루어보도록 하겠습니다. 

 

우선, 아래와 같은 5개의 간단한 문서가 있다고 가정해보겠습니다.

d1 = '사과 바나나 오렌지 바나나 멜론'

d2 = '포도 포도 오렌지 오렌지 수박'

d3 = '사과 사과 바나나 오렌지 수박'

d4 = '포도 사과 멜론 딸기 딸기'

d5 = '딸기 수박 멜론 바나나 오렌지'

이 d1~d5를 기준으로 각 문서별로 가장 유사한 document를

단어 기준 TF-IDF 임베딩 및 코사인 유사도로 찾아보도록 하겠습니다.

 

 

파이썬 TF-IDF 벡터화

사이킷런의 TfidfVectorizer 함수를 이용하여 TF-IDF 임베딩을 쉽게 진행할 수 있습니다.

 

벡터화 결과를 쉽게 알아볼 수 있도록 단어 목록 및 문서 목록을 column과 index로 각각 가지는

데이터프레임 형태로 시각화한 예시 코드는 아래와 같습니다.

from sklearn.feature_extraction.text import TfidfVectorizer
import pandas as pd

documents = [d1, d2, d3, d4, d5]

# TF-IDF 점수 변환 모델 준비
tf_idf_model = TfidfVectorizer().fit(documents)

# 단어 이름 리스트 순서대로 반환(0번으로 지정된 토큰부터)
word_id_list = sorted(tf_idf_model.vocabulary_.items(), key=lambda x: x[1], reverse=False)
word_list = [x[0] for x in word_id_list]

# 용이한 시각화를 위하여 데이터프레임 변환
tf_idf_df = pd.DataFrame(tf_idf_model.transform(documents).toarray(), columns = word_list, index = ['d1', 'd2', 'd3', 'd4', 'd5'])

tf_idf_df

각 문서별로 토큰(단어)에 대한 TF-IDF 수치가 잘 요약된 모습을 확인할 수 있었습니다.

반응형

코사인 유사도(Cosine Similarity) 기반 유사 문서 탐색 과정

위에서 만든 문서별 TF-IDF 벡터끼리의 코사인 유사도를 통하여

문서 간 유사도를 계산하고, 가장 비슷한 문서가 어떤 문서인지를 쉽게 찾을 수 있습니다.

 

위에서 생성한 데이터프레임을 사이킷런의 cosine_similarity 함수에

대입하여 코사인 유사도를 바로 계산할 수 있습니다.

from sklearn.metrics.pairwise import cosine_similarity

cos_sim_df = pd.DataFrame(cosine_similarity(tf_idf_df, tf_idf_df), columns = ['d1', 'd2', 'd3', 'd4', 'd5'], index = ['d1', 'd2', 'd3', 'd4', 'd5'])

cos_sim_df

모든 문서 pair 간의 코사인 유사도가 위의 데이터프레임 형태로 요약되었습니다.

 

위의 유사도를 기반으로 각 document별 가장 유사한 문서를 찾아보도록 하겠습니다.

자기 자신 문서와의 유사도는 항상 1인데, 이는 제외해야하므로 조건 인덱싱 후

최대값을 가지는 인덱스를 idxmax 함수를 통하여 가져왔습니다.

cos_sim_df[cos_sim_df < (1-1e-6)].idxmax()

# 결과
d1    d3
d2    d5
d3    d1
d4    d5
d5    d1
dtype: object

각 문서별로 가장 비슷하다고 인식된 문서의 목록을 쉽게 찾을 수 있었습니다.

 

이 목록은 양방향의 관계가 성립하지 않는 점을 알 수 있습니다.(d2 -> d5 but d5 -> d2가 아님)