「intfloat/multilingual-e5-large」を使ってとりあえずベクトル検索

概要

ベクトル検索は、ベクトルDBを構築して検索をかけるのが主流で、オープンソースでベクトルDBを構築する手順を以下の記事に記載している。

今回はベクトルDBすら構築するのが煩わしい方向けに、オープンソース界で優れた結果を残したエンベディングモデルである「intfloat/multilingual-e5-large」というモデルを使って、簡単にベクトル検索を体験するプログラムを作成する。

ただし、今回の環境はRDBのようにテーブル内のデータをすべて参照して検索を行うような挙動をするため、ベクトル検索の件数が増えても性能が劣化しない利点は得られない環境となっている。
また、プログラム自体は、CUDAを利用したプログラムになっているが、CPUのみの環境でも動作すると想定している。

環境

OS:Windows 11
GPU:GeForce RTX 4090
CPU:i9-13900KF
memory:64G
python:3.10.10
pytorch:2.0.1
CUDA:11.8
cuDNN:8.8

以下の環境でも動作確認済み。
GPU:GeForce RTX 3060 laptop
CPU:i7-10750H
memory:16G

環境構築

以下のコマンドでライブラリをインストールする。

CUDAなしの場合

pip install numpy
pip install transformers
pip install torch
pip install scipy

CUDAありの場合

pip install numpy
pip install transformers
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
pip install scipy

プログラム

ベクトル検索準備

以下のプログラムを実行すると、変数.text_list内のquestionテキストをエンベディング（ベクトル化）し、CSVに出力する。

import csv
from transformers import AutoTokenizer, AutoModel
import torch
from torch import Tensor
import torch.nn.functional as F

model_name = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# CUDAが利用可能かチェックし、利用可能であればデバイスをCUDAに設定
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# モデルをデバイスに移動
model = AutoModel.from_pretrained(model_name).to(device)

def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

text_list = [
    {"question":"query: あなたの名前は？", "answer":"ネムル"},
    {"question":"query: 好きなものは？", "answer":"寝ること"},
    {"question":"query: 好き食べ物は？", "answer":"芋"},
    {"question":"query: 元気がでるのは？", "answer":"寝ること"},
    {"question":"query: 最近うれしかったことは？", "answer":"今日晴れてたこと"},
    {"question":"query: 歳は？", "answer":"知らない"},
    {"question":"query: 出身は？", "answer":"ヨネズのコロニー"},
    {"question":"query: 好きな人は？", "answer":"みはる"},
    {"question":"query: 嫌いなことは？", "answer":"頑張ること"},
    {"question":"query: 得意な勉強は？", "answer":"勉強とかできない"},
    {"question":"query: 朝型か夜型か？", "answer":"特にない"},
    {"question":"query: 普段、何をしていますか？", "answer":"だいたい寝てるかも"},
    {"question":"query: 趣味は？", "answer":"寝る以外とかにない"},
    {"question":"query: 最近読んだ本や漫画は？", "answer":"本？"},
    {"question":"query: 印象に残っている旅行は？", "answer":"シティより遠くに行ったことない？"},
    {"question":"query: 最近挑戦したことは？", "answer":"しゃべること"},
    {"question":"query: 最近感動したことは？", "answer":"最近暖かくて寝心地がいい"},
    {"question":"query: 最近見た映画やドラマは？", "answer":"映画？"},
]

# CSVファイルを開き、書き込みます
with open('q_and_a_with_vectors.csv', mode='w', encoding='utf-8', newline='') as csv_file:
    fieldnames = ['id', 'question', 'answer', 'vector']
    writer = csv.DictWriter(csv_file, fieldnames=fieldnames)
    
    writer.writeheader()
    id = 1
    for text in text_list:
        inputs = tokenizer(text['question'], return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
        
        with torch.no_grad():  # 勾配計算を不要にする
            outputs = model(**inputs)
        
        embeddings = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
        embeddings = F.normalize(embeddings, p=2, dim=1)
        vector_string = ",".join([f"{x:.20f}" for x in embeddings[0].cpu().numpy()])  # ベクトルを文字列に変換
        
        writer.writerow({'id': id, 'question': text['question'], 'answer': text['answer'], 'vector': vector_string})
        id += 1

CSVへは、（一部省略しているが）以下のような出力結果となる。
ここで、ベクトル化されているのは、questionのテキストで、questionをベクトル検索することで、類似する質問とその答えを取得する。

id,question,answer,vector
1,query: あなたの名前は？,ネムル,"0.03061610460281372070,-0.01759614422917366028,-0.02009218186140060425, ..."
2,query: 好きなものは？,寝ること,"0.04883756116032600403,-0.01367216464132070541,-0.01473792176693677902, ..."
3,query: 好き食べ物は？,芋,"0.03631471097469329834,-0.02889459580183029175,-0.00063249614322558045, ..."

ベクトル検索実行

今回は「出身地は？」という質問を投げて、この質問と類似度の高い質問をCSVから取得し、そのanswerと共に出力する。

import csv
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F
from torch import Tensor
from scipy.spatial.distance import cdist

model_name = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def average_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

def cosine_similarity(v1, v2):
    return 1 - cdist([v1], [v2], 'cosine')[0][0]

# クエリテキストをエンベディング
query_text = "query: 出身地は？"
inputs = tokenizer(query_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
with torch.no_grad():
    outputs = model(**inputs)

query_embeddings = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
query_embeddings = F.normalize(query_embeddings, p=2, dim=1).cpu().numpy()[0]

# CSVファイルを読み込み、各レコードとクエリの類似度を計算
similarities = []

with open('q_and_a_with_vectors.csv', mode='r', encoding='utf-8') as csv_file:
    reader = csv.DictReader(csv_file)
    for row in reader:
        vector = np.array([float(x) for x in row['vector'].split(',')])
        similarity = cosine_similarity(query_embeddings, vector)
        similarities.append((row, similarity))

# 類似度でソートし、上位3つの結果を取得
top_matches = sorted(similarities, key=lambda x: x[1], reverse=True)[:3]

for i, (row, similarity) in enumerate(top_matches, 1):
    print(f"Match #{i}:")
    print(f"  Question: {row['question']}")
    print(f"  Answer: {row['answer']}")
    print(f"  Similarity: {similarity:.20f}\n")

上記のベクトル検索を実行すると、結果は以下となる、（ベクトル検索の精度を考慮して、類似度が高い３つのテキストを出力する。）

Match #1:
  Question: query: 出身は？
  Answer: ヨネズのコロニー
  Similarity: 0.98094844793202040645

Match #2:
  Question: query: あなたの名前は？
  Answer: ネムル
  Similarity: 0.87477910777269485276

Match #3:
  Question: query: 歳は？
  Answer: 知らない
  Similarity: 0.86065601455392459762

今回のテキストは「出身地は？」と「出身は？」で、明らかに近いテキストということもあるが、類似度が「0.98094844793202040645」となり、想定通りの結果となった。

補足

「intfloat/multilingual-e5-large」を使う時に、入力テキストの接頭辞に「query: 」や「passage: 」というプレフィックスをつけることが推奨されている。

「query: 」は質問や検索キーワード、または比較したいテキストを示す時に使い、「passage: 」はその質問やキーワードに対する情報や答えが含まれるテキストを示す時に使う。

「query: 」の方が汎用的らしいので、迷ったら「query: 」を使った方がいいらしい。