Python で maximal marginal relevance (MMR)
maximal marginal relevance (MMR) を実装したことなかったので雑に実装してみる。大体GPT4が書いたのだけど。
code:python
import numpy as np
def MMR(
doc_embeddings: np.ndarray, query_embedding: np.ndarray, lambda_param=0.5, top_n=100
):
# 関連性スコアを計算(ベクトル化)
query_norm = np.linalg.norm(query_embedding)
doc_norms = np.linalg.norm(doc_embeddings, axis=1)
relevance_scores = np.dot(doc_embeddings, query_embedding) / (
doc_norms * query_norm
)
# すべてのドキュメント間の類似性スコアを計算(ベクトル化)
similarity_scores = np.dot(doc_embeddings, doc_embeddings.T) / np.outer(
doc_norms, doc_norms
)
# MMRスコアを格納する配列と選択されたドキュメントのリストを初期化
selected_docs = np.zeros(top_n, dtype=int)
mmr_scores = np.zeros(top_n)
for i in range(top_n):
if i == 0: # 最初は関連性スコアのみを使用
next_selected = np.argmax(relevance_scores)
mmr_values = relevance_scores.copy() # 初期化
else:
# 既に選択されたドキュメントとの最大類似性を計算
max_similarities = similarity_scores[selected_docs:i, :].max(axis=0) # MMRスコアを計算(ベクトル化)
mmr_values = (
lambda_param * relevance_scores - (1 - lambda_param) * max_similarities
)
mmr_values[selected_docs:i] = -np.inf # 既に選択されたドキュメントは除外 next_selected = np.argmax(mmr_values)
selected_docsi = next_selected return selected_docs, mmr_scores
query_embedding = np.random.rand(384)
doc_embeddings = np.random.rand(300, 384)
selected_docs, scores = MMR(doc_embeddings, query_embedding, lambda_param=0.5, top_n=10)