埋め込みモデルe5-mistral-7b-instructを使ってみた。

概要

埋め込みモデルe5-mistral-7b-instructを使って、テキスト間のコサイン類似度を計算する方法をまとめる。
※内容が間違っている可能性があります、ご容赦ください。

e5-mistral-7b-instruct

今回使用する、埋め込みモデルです。
2024年2月17日現在、MTEB LeaderboardのEnglishで4位となっています。
多言語でも使えますが、英語の使用が勧められています。

huggingface.co

実装

こちらを参考に実装しています。

huggingface.co

import torch
import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

def last_token_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]

documents = [
    "次の日は良い天気です",
    "明日の天気は晴れです",
    "明日の天気は雨です",
    "晴れが好きです",
    "私はリンゴが好きです"
]

tokenizer = AutoTokenizer.from_pretrained('intfloat/e5-mistral-7b-instruct')
model = AutoModel.from_pretrained('intfloat/e5-mistral-7b-instruct')

max_length = 4096

# 入力テキストをトークン化する
batch_dict = tokenizer(documents, max_length=max_length - 1, return_attention_mask=False, padding=False, truncation=True)

# eos_token_idをすべてのinput_idsに追加する
batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
batch_dict = tokenizer.pad(batch_dict, padding=True, return_attention_mask=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# 埋め込みを正規化する
embeddings = F.normalize(embeddings, p=2, dim=1)
embeddings_list = embeddings.tolist()

def similarity(query_embedding, document_embedding):

    query_embedding = np.array(query_embedding)
    document_embedding = np.array(document_embedding)
    
    # NumPy配列を2次元配列に変換
    query_embedding = query_embedding.reshape(1, -1)
    document_embedding = document_embedding.reshape(1, -1)
    
    result = cosine_similarity(query_embedding, document_embedding)
    
    return result

cos_similarity=[]
for i in range(len(embeddings)-1):
    cos_similarity.append(similarity(embeddings_list[0], embeddings_list[i+1]))

print(f"query={documents[0]}")
for text, score in zip(documents[1:], cos_similarity):
    print(f"text={text}, score={score[0][0]}")

query=次の日は良い天気です
text=明日の天気は晴れです, score=0.89586343422919
text=明日の天気は雨です, score=0.8419866323866365
text=晴れが好きです, score=0.8280332788608935
text=私はリンゴが好きです, score=0.6293116246996098

結果

コサイン類似度を計算すると、「次の日は良い天気です」に対して、「明日の天気は晴れです」が高いスコアとなっていました。

疑問点

今回はクエリと指示の部分は削除して使用していますが、この使用方法が正しいのかわかっていません。