news_summarizer / cosine_similarity_calculator.py
Kota Takahashi
説明文を修正
44bec05
import gensim
from sklearn.metrics.pairwise import cosine_similarity
class CosineSimilarityCalculator:
model_path = 'ja/ja.bin'
def __init__(self):
"""
CosineSimilarityCalculatorクラスを初期化し、
事前トレーニング済みのWord2Vecモデルをロード
Parameters:
- なし
Returns:
- なし
"""
self.model = gensim.models.Word2Vec.load(CosineSimilarityCalculator.model_path)
def _convert_to_2d_array(self, vector):
"""
埋め込みベクトルを2次元配列に変換
Parameters:
- vector (numpy.ndarray): 変換する1次元配列のベクトル
Returns:
- vector_2d (numpy.ndarray): 変換後の2次元配列のベクトル
"""
return vector.reshape(1, -1)
def _calculate_cosine_similarity(self, embedding1, embedding2):
"""
cos類似度を計算
Parameters:
- embedding1 (numpy.ndarray): 1つ目の単語ベクトル(2次元配列)
- embedding2 (numpy.ndarray): 2つ目の単語ベクトル(2次元配列)
Returns:
- similarity (numpy.ndarray): cos類似度
"""
return cosine_similarity(embedding1, embedding2)
def calculate_similarity(self, search_word, article_keyword_list):
"""
指定された検索ワードと記事のキーワードリストの間のcos類似度を計算
モデルにない単語の場合はエラーメッセージを出力しブレイクする
Parameters:
- search_word (str): 検索ワード
- article_keyword_list (list): 記事のキーワードリスト
Returns:
- similarities (dict): 記事キーワードとそれぞれの検索ワードのcos類似度を含むdictを作成
モデルにない単語の場合はNoneを返す
"""
# 検索ワードの埋め込みベクトルを取得
if search_word in self.model.wv:
search_embedding = self.model.wv[search_word]
else:
return None
similarities = {}
# 記事キーワードの埋め込みベクトルを取得し、cos類似度を計算
for keyword in article_keyword_list:
if keyword in self.model.wv:
keyword_embedding = self.model.wv[keyword]
search_embedding_2d = self._convert_to_2d_array(search_embedding)
keyword_embedding_2d = self._convert_to_2d_array(keyword_embedding)
similarity = self._calculate_cosine_similarity(search_embedding_2d, keyword_embedding_2d)
similarities[keyword] = similarity[0][0]
else:
similarities[keyword] = None
return similarities