概要

以前、ローカルでベクトル検索を実施することに関して、とりあえず試す手順を作成した。
「intfloat/multilingual-e5-large」モデルの性能がある程度高く、CPUのみの環境でも優れた結果を残せることが分かっている。

今回はこの手順の派生で、ローカルで構築したベクトル検索環境を、簡単な手順で他のユーザに共有可能な手順を作成する。
RDBを利用するため、ベクトル検索の特徴である件数増加による性能の劣化がないというメリットは得られません。

構成としては以下の構成を想定している。
「setup_vector_db.bat」を実行することで、環境を構築し、「start_vector_db.bat」で検索用のアプリケーションを起動する。

動画でも解説しています。

環境

OS:Windows 11
GPU:GeForce RTX 3060 laptop　※GPUは不要
CPU:i7-10750H
memory:16G
python:3.10.11

環境構築

以下のコマンドを実行し、pythonの仮想環境を構築する。（おすそ分けを受けるユーザも必須の手順）

python -m venv vector_sqlite_venv
.\vector_sqlite_venv\Scripts\activate
pip install numpy
pip install transformers
pip install torch
pip install scipy
pip install flask
pip install flask-cors
pip install requests
deactivate

以下のフォルダ構成となるように、source配下にファイルを作成してください。

vector_sqlite_venv/
├── source/
│   ├── setup_vector_db.bat
│   ├── start_vector_db.bat
│   ├── text_list.csv
│   ├── create_vector_table.py
│   ├── insert_vector_table.py
│   ├── select_vector_table.py
│   ├── search_vector_table.py
│   └── index.html
└── Scripts/
    └── activate

setup_vector_db.bat

@echo off
call .\..\Scripts\activate
python create_vector_table.py
python insert_vector_table.py
pause

start_vector_db.bat

@echo off
call .\..\Scripts\activate
start "Python Script" python search_vector_table.py
timeout /t 5 /nobreak > nul
start http://127.0.0.1:5000

text_list.csv

id,question,answer
1,query: あなたの名前は？ネムル
2,query: 好きなものは？,寝ること
3,query: 好き食べ物は？,芋
4,query: 元気がでるのは？,寝ること
5,query: 最近うれしかったことは？,今日晴れてたこと
6,query: 歳は？,知らない
7,query: 出身は？,ヨネズのコロニー
8,query: 好きな人は？,みはる
9,query: 嫌いなことは？,頑張ること
10,query: 得意な勉強は？,勉強とかできない
11,query: 朝型か夜型か？,特にない
12,query: 普段、何をしていますか？,だいたい寝てるかも
13,query: 趣味は？,寝る以外特にない
14,query: 最近読んだ本や漫画は？,本？
15,query: 印象に残っている旅行は？,シティより遠くに行ったことない？
16,query: 最近挑戦したことは？,しゃべること
17,query: 最近感動したことは？,最近暖かくて寝心地がいい
18,query: 最近見たドラマや映画は？,映画？
19,query: 運動は何が得意？,運動は苦手
20,query: 好き飲み物は？,雨水は嫌いだけど、水
21,query: 季節の中で好きなのは？,春が好き
22,query: お気に入りの音楽は？,あんまり知らない
23,query: 誰と一緒に住んでいますか？,最近はみはる
24,query: テレビ番組はよく見ますか？,テレビって何？
25,query: 異性の好みは？,異性って誰の事？
26,query: お気に入りの映画ジャンルは？,映画って何？
27,query: 空いた時間はどう過ごしますか？,寝てる
28,query: 今の生活で満足していますか？,特に不満はない
29,query: 将来の夢は？,ずっと寝ること
30,query: 自分を色で表すと？,茶色
31,query: 一番長く話したことは何ですか？,あんまり話せない
32,query: 友達は多いですか？,少ない
33,query: どんな映画が嫌い？,映画が何かわからない
34,query: 絵を描くことは好きですか？,絵はあまり得意じゃない
35,query: 好きな季節のイベントは？,晴れてる春
36,query: 休日の過ごし方は？,寝てる
37,query: 何色の服をよく着ますか？,茶色しかない
38,query: 最後に笑ったことは？,みはるが川でひっくり返ったとき
39,query: 携帯電話はよく使いますか？,携帯って何かわからない
40,query: 自分の長所と短所は？,一人じゃ何もできない
41,query: 自分を動物に例えると？,ダンゴムシ
42,query: 好きなスポーツは？,スポーツが何かわからないよ
43,query: 自炊はしますか？,まったくしない
44,query: 好きなお菓子は？,虫の幼虫
45,query: 最近買ったものは？,買うことなんてない
46,query: 好きなアニメはありますか？,アニメ？
47,query: パソコンは使いますか？,みはるがよく使ってるけど、ネムルはよくわからない
48,query: お気に入りのレストランは？,レストランとかはない
49,query: 自分の性格を一言で表すと？,寝坊助
50,query: 最近感じた小さな幸せは？,昼まで寝れた時

create_vector_table.py

import sqlite3

# データベースファイルを作成し接続する
conn = sqlite3.connect('vector_db.sqlite')
c = conn.cursor()

# ベクトルデータを保存するためのテーブルを作成
c.execute('''
CREATE TABLE qa_vectors_table (
    id INTEGER PRIMARY KEY,
    question TEXT,
    answer TEXT,
    vector BLOB  -- ベクトルデータをBLOBとして保存
)
''')

conn.commit()
conn.close()

insert_vector_table.py

import csv
import sqlite3
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)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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]

# データベースへの接続
conn = sqlite3.connect('vector_db.sqlite')
c = conn.cursor()

# CSVファイルから質問と回答を読み込む
with open('text_list.csv', mode='r', encoding='utf-8') as csv_file:
    reader = csv.DictReader(csv_file)
    for row in reader:
        question = f"{row['question']}"
        answer = row['answer']
        inputs = tokenizer(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_blob = sqlite3.Binary(embeddings.numpy().tobytes())

        # データベースにレコードを挿入
        c.execute('INSERT INTO qa_vectors_table (question, answer, vector) VALUES (?, ?, ?)', (question, answer, vector_blob))

# 変更をコミットし、データベース接続を閉じる
conn.commit()
conn.close()

select_vector_table.py

import sqlite3
from transformers import AutoTokenizer, AutoModel
import torch
import numpy as np
from torch import Tensor
import torch.nn.functional as F

# モデルの準備
model_name = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
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]

def cosine_similarity(a, b):
    # ベクトル a と b の形状を確認し、1次元になっていることを保証
    a = a.flatten()
    b = b.flatten()
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# ユーザーのクエリを入力
query = input("Please enter your query: ")

# クエリをベクトルに変換
inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
with torch.no_grad():
    outputs = model(**inputs)
query_vec = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
query_vec = F.normalize(query_vec, p=2, dim=1)
query_vec = query_vec.cpu().numpy()

# データベース接続
conn = sqlite3.connect('vector_db.sqlite')
c = conn.cursor()

# データベースからすべてのベクトルを取得
c.execute("SELECT id, question, answer, vector FROM qa_vectors_table")
results = c.fetchall()

# 類似度計算
max_similarity = -1
best_match = None
for row in results:
    db_id, db_question, db_answer, db_vector = row
    db_vector = np.frombuffer(db_vector, dtype=np.float32)  # BLOBをnumpy配列に変換
    similarity = cosine_similarity(query_vec, db_vector)
    if similarity > max_similarity:
        max_similarity = similarity
        best_match = (db_question, db_answer, similarity)

# 結果の表示
if best_match:
    print(f"Most similar question: {best_match[0]}")
    print(f"Answer: {best_match[1]}")
    # Similarity score を正しく表示
    print(f"Similarity score: {float(best_match[2]):.4f}")
else:
    print("No matching question found.")

# データベース接続を閉じる
conn.close()

search_vector_table.py

from flask import Flask, request, jsonify
from flask_cors import CORS
import sqlite3
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModel
import torch.nn.functional as F

app = Flask(__name__)
CORS(app)

model_name = "intfloat/multilingual-e5-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModel.from_pretrained(model_name).to(device)

def average_pool(last_hidden_states, attention_mask):
    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(a, b):
    a = a.flatten()
    b = b.flatten()
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

@app.route('/')
def index():
    # ファイルをUTF-8エンコーディングで開く
    with open('index.html', 'r', encoding='utf-8') as file:
        return file.read()

@app.route('/search', methods=['POST'])
def search():
    query = request.json['query']
    inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    query_vec = average_pool(outputs.last_hidden_state, inputs['attention_mask'])
    query_vec = F.normalize(query_vec, p=2, dim=1)
    query_vec = query_vec.cpu().numpy()

    conn = sqlite3.connect('vector_db.sqlite')
    c = conn.cursor()
    c.execute("SELECT id, question, answer, vector FROM qa_vectors_table")
    results = c.fetchall()
    conn.close()

    max_similarity = -1
    best_match = None
    for row in results:
        db_id, db_question, db_answer, db_vector = row
        db_vector = np.frombuffer(db_vector, dtype=np.float32)
        similarity = cosine_similarity(query_vec, db_vector)
        if similarity > max_similarity:
            max_similarity = similarity
            best_match = {'question': db_question, 'answer': db_answer, 'similarity': float(similarity)}

    return jsonify(best_match)

if __name__ == '__main__':
    app.run(debug=True)

index.html

<!DOCTYPE html>
<html lang="ja">
<head>
    <meta charset="UTF-8">
    <title>ベクトル検索インターフェース</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            padding: 20px;
            margin: 0;
            background-color: #f4f4f4;
        }
        #queryInput {
            width: 80%;
            padding: 10px;
            font-size: 16px;
        }
        #searchButton {
            padding: 10px 20px;
            font-size: 16px;
            cursor: pointer;
        }
        #result {
            margin-top: 20px;
            background-color: white;
            padding: 20px;
            border-radius: 5px;
        }
        #loadingText {
            display: none;
        }
    </style>
</head>
<body>
    <input type="text" id="queryInput" placeholder="クエリを入力してください">
    <button id="searchButton" onclick="submitQuery()">検索</button>
    <p id="loadingText">検索中...</p>
    <div id="result">
        <p id="question"></p>
        <p id="answer"></p>
        <p id="similarity"></p>
    </div>

    <script>
    function submitQuery() {
        const query = document.getElementById('queryInput').value;
        const loadingText = document.getElementById('loadingText');
        const searchButton = document.getElementById('searchButton');
        searchButton.disabled = true;
        loadingText.style.display = 'block';

        fetch('http://localhost:5000/search', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json'
            },
            body: JSON.stringify({ query: query })
        })
        .then(response => response.json())
        .then(data => {
            if(data) {
                document.getElementById('question').textContent = "質問: " + data.question;
                document.getElementById('answer').textContent = "回答: " + data.answer;
                document.getElementById('similarity').textContent = "類似度スコア: " + data.similarity.toFixed(4);
            } else {
                document.getElementById('question').textContent = "該当する質問が見つかりませんでした。";
                document.getElementById('answer').textContent = "";
                document.getElementById('similarity').textContent = "";
            }
            loadingText.style.display = 'none';
            searchButton.disabled = false;
        })
        .catch(error => {
            console.error('エラー:', error);
            document.getElementById('question').textContent = "エラーが発生しました。";
            document.getElementById('answer').textContent = "";
            document.getElementById('similarity').textContent = "";
            loadingText.style.display = 'none';
            searchButton.disabled = false;
        });
    }
    </script>
</body>
</html>

ベクトルDB構築

【「setup_vector_db.bat」を実行する】 or 【別ユーザが作成した「vector_db.sqlite」をコピーする】でDBを構築する。

※検索するデータを変更したい場合は、「text_list.csv」のデータを改変し、「vector_db.sqlite」を削除後、再度「setup_vector_db.bat」を実行する。

ベクトル検索実施

「start_vector_db.bat」を実行すると、ブラウザで以下のような画面が表示される。

検索を行うと以下のような結果となる。

感想

利用しているエンベディングモデルについては、トークン数が512と非常に少ないが、ベクトル検索を行うためのベースのモデルとしては、やはり十分な性能があると感じている。

今回構築した環境については、ベクトルDBではなく、RDBを利用しているため、件数次第では、性能劣化が起きる。環境構築が簡単なベクトルDBがリリースされれば、そちらを利用することをオススメする。