WeaviateとOSSエンベディングモデルを使ったベクトルDB構築

記事の概要

RAGの技術を支えているベクトル検索を体験したいと思いいたったが、Azure AI SearchやOpenAIのエンベディングモデルなどの課金サービスが一般的となっているため、「とりあえず試したい」と思った時に敷居が高いというのが現状だと感じている。

そこで、オープンソースのベクトルDB（Weaviate）と日本語対応のエンベディングモデル（BERT）を使った、ベクトル検索を体験する環境を構築する。

※コサイン類似度やベクトル、エンベディングの考え方については、省略します。
※Weaviateはv4もリリースされているみたいだが、今回はv3で構築している。

動画でも紹介してます。

環境

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

環境構築

環境構築では、以下のフローチャートの「PCでこの手順を実施」を実施する。
手順を実施すると、フローチャートの「Docker」内の構成が作成され、簡単なベクトル検索を実施できる。

ステップ1: Dockerを使ってWeaviateをインストール

1. Docker Desktopのインストール

2. Weaviateのコンテナを起動

PowerShellまたはコマンドプロンプトを開き、以下のコマンドを実行してWeaviateの最新版を起動する。

docker run -d --name weaviate -p 8080:8080 semitechnologies/weaviate:latest

これにより、Weaviateがバックグラウンドで実行され、8080ポートでアクセス可能になる。

3.Weaviateの動作確認

docker desktopを開き、キャプチャのURLを開くと、以下のjsonがブラウザのレスポンスとして表示されることを確認する。

{"links":[{"href":"/v1/meta","name":"Meta information about this instance/cluster"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/schema","href":"/v1/schema","name":"view complete schema"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/schema","href":"/v1/schema{/:className}","name":"CRUD schema"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/objects","href":"/v1/objects{/:id}","name":"CRUD objects"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/classification,https://weaviate.io/developers/weaviate/api/rest/classification#knn-classification","href":"/v1/classifications{/:id}","name":"trigger and view status of classifications"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/well-known#liveness","href":"/v1/.well-known/live","name":"check if Weaviate is live (returns 200 on GET when live)"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/well-known#readiness","href":"/v1/.well-known/ready","name":"check if Weaviate is ready (returns 200 on GET when ready)"},{"documentationHref":"https://weaviate.io/developers/weaviate/api/rest/well-known#openid-configuration","href":"/v1/.well-known/openid-configuration","name":"view link to openid configuration (returns 404 on GET if no openid is configured)"}]}

ステップ2: 日本語対応のエンベディングモデルの導入

Pythonと必要なライブラリ（特にHugging FaceのTransformers）を使用して、日本語のテキストデータからエンベディングを生成する。

1. python仮想環境を構築

python -m venv embedding_env

2. 仮想環境をactivate

.\embedding_env\Scripts\activate

3. 必要なライブラリをインストール

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
pip install transformers
pip install weaviate-client
pip install fugashi ipadic

ステップ3: テキストデータのエンベディング生成

1. テキストデータからエンベディングを生成

日本語対応のエンベディングモデル（例：cl-tohoku/bert-base-japanese）を使用して、テキストデータからエンベディングを生成する。

# 必要なライブラリをインポートします。
from transformers import BertModel, BertJapaneseTokenizer
import torch

# 使用するモデルの名称を指定します。ここでは、日本語対応のBERTモデル「cl-tohoku/bert-base-japanese」を使用します。
model_name = "cl-tohoku/bert-base-japanese"

# 指定したモデル名で、トークナイザー（テキストをトークンに分割するためのツール）を初期化します。
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)

# 指定したモデル名で、BERTモデルを初期化します。
model = BertModel.from_pretrained(model_name)

# エンベディングを生成したいテキストデータです。
text = "あなたのテキストをここに入れてください。"

# トークナイザーを使用して、テキストデータをモデルが理解できる形式に変換します。
# return_tensors="pt"は、PyTorchテンソルを出力するよう指定します。
# padding=Trueは、入力の長さを揃えるためにパディングを行うことを指定します。
# truncation=Trueは、指定された最大長さを超える入力を切り詰めることを指定します。
# max_length=128は、入力の最大長さを指定します。
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)

# モデルを使用して、入力からエンベディングを生成します。
outputs = model(**inputs)

# 最後の隠れ層の状態から平均を取り、テキストのエンベディングを取得します。
# .mean(dim=1)は、トークンごとのエンベディングの平均を計算します。
# .detach()は、計算グラフからこのテンソルを分離して、以降の計算で勾配が不要であることを示します。
# .numpy()は、テンソルをNumPy配列に変換します。
embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()

ステップ4: Weaviateにテキストエンベディングをインサート

1.生成したエンベディングとテキストデータをWeaviateにインサート

以下のPythonスクリプトを使用して、生成したエンベディングとテキストデータをWeaviateにインサートする。
今回の構成では、「"vectorizer": "none"」としているのは重要で、インサート前にベクトル化するための設定。
vectorizerに、エンベディングモデルを指定すると、テキストデータを登録すると自動で、vector項目にベクトル化される処理フローとなる。

# WeaviateのPythonクライアントをインポートします。
import weaviate

# Weaviateサーバーへの接続を初期化します。ここではローカルホスト上の標準ポート(8080)を指定しています。
client = weaviate.Client("http://localhost:8080")

# Weaviateのスキーマを定義します。このスキーマには、TextDataクラスが含まれています。
schema = {
    "classes": [{
        "class": "TextData", # クラス名を定義します。
        "vectorizer": "none",
        "properties": [{ # クラスのプロパティを定義します。
            "name": "text", # テキストデータを保持するプロパティです。
            "dataType": ["string"], # データタイプは文字列です。
            "description": "テキストデータ"
        }]
    }]
}

# 定義したスキーマをWeaviateに作成します。
client.schema.create(schema)

# 定義したデータオブジェクトをWeaviateのTextDataクラスにインサートします。
client.batch.configure(batch_size=1)  # Configure batch
with client.batch as batch:
    properties = {
        "text": text,
    }
    batch.add_data_object(properties, "TextData", vector=embeddings.tolist()[0])

今回は「text」という項目を持つ、TextDataというテーブル（クラス）を作成している。

ステップ5: セマンティック検索の実行

1. Weaviateの検索機能を使ったセマンティック検索を実行

Weaviateの検索機能（with_near_vector）を使用して、インサートされたテキストデータに対してセマンティック検索を実行する。

text2 = "これで検索します。"
inputs2 = tokenizer(text2, return_tensors="pt", padding=True, truncation=True, max_length=128)
outputs2 = model(**inputs2)
embeddings2 = outputs2.last_hidden_state.mean(dim=1).detach().numpy()

# Weaviateのクエリ機能を使用して、特定のエンベディングに最も近いテキストデータを検索します。
query_result = client.query.get("TextData", ["text"]).with_near_vector({"vector": embeddings2.tolist()[0]}).with_limit(2).with_additional(['certainty']).do()

# 検索結果を出力します。この結果には、クエリに最も近いテキストデータが含まれます。
print(query_result)

これらのステップに従って、Windows PC上にWeaviateをインストールし、日本語のテキストデータからエンベディングを生成してWeaviateにインサートし、セマンティック検索を実行することができる。

複数データからの検索（ついでにCUDAでエンベディング）

上記で構築したクラスとは別のクラスを作成し、そこに複数のデータを登録して、ベクトル検索を行う。

テキストデータからエンベディングを生成・インサート

from transformers import BertModel, BertJapaneseTokenizer
import torch
import weaviate
from datetime import datetime, timezone

model_name = "cl-tohoku/bert-base-japanese"
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)

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

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

client = weaviate.Client("http://localhost:8080")

class_obj = {
    "class": "AddColumnTextData",
    "vectorizer": "none",
    "description": "テキストデータとそのエンベディング、作成日時を含むクラス",
    "properties": [
        {
            "name": "unique_id",
            "dataType": ["string"],
            "description": "ユニークなID"
        },
        {
            "name": "text",
            "dataType": ["string"],
            "description": "テキストデータ"
        },
        {
            "name": "create_date",
            "dataType": ["date"],
            "description": "作成日時"
        }
    ]
}


client.schema.create_class(class_obj)

text_list = [
 "sd-webui-animatediffで高解像度・高フレームレートのクロマキーの動画を作成（v3_sd15_mm.ckpt）",
 "sd-webui-animatediffで高解像度・高フレームレートのクロマキーの動画を作成",
 "M1 MACで深層強化学習（スペック的には厳しい）",
 "画像反転・高解像度化・GIF化・MP4化　pythonプログラム",
 "SDXLモデルのDreambooth",
 "GPT×マーメイド記法で図やフローを出力",
 "DETR（End-to-End Object Detection with Transformers）を使った動画での物体検知",
 "DETR（End-to-End Object Detection with Transformers）物体検知（自前データでのファインチューニング）",
 "動画からOpenpose用の骨格データをPNGで出力",
 "DQN（Deep Q-learning Network）のハイパーパラメータ更新",
 "DETR（End-to-End Object Detection with Transformers）物体検知",
 "WindowsでのDQN（Deep Q-learning Network）",
 "Google ColabでのDQN（Deep Q-learning Network）"
]

id = 1
for text in text_list:
    # 入力をトークナイズし、デバイスに移動
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
    
    # モデルを実行してエンベディングを取得
    with torch.no_grad():  # 勾配計算を不要にする
        outputs = model(**inputs)
    
    embeddings = outputs.last_hidden_state.mean(dim=1).detach().cpu().numpy()  # 結果をCPUに戻す
    
    data_object = {
        "unique_id": str(id),
        "text": text,
        "create_date": datetime.now(timezone.utc).isoformat()
    }
    client.data_object.create(data_object, "AddColumnTextData", vector=embeddings.tolist()[0])
    id += 1

３つの項目を持つ、AddColumnTextDataというテーブル（クラス）を作成している。

インサートしたデータの確認

import weaviate

client = weaviate.Client("http://localhost:8080")

query = """
{
  Get {
    AddColumnTextData {
      unique_id
      text
      create_date
    }
  }
}
"""

result = client.query.raw(query)
print(result)

以下の実行結果が得られる。

{'data': {'Get': {'AddColumnTextData': [{'create_date': '2024-03-05T14:07:56.075299Z', 'text': '画像反転・高解像度化・GIF化・MP4化\u3000pythonプログラム', 'unique_id': '4'}, {'create_date': '2024-03-05T14:07:56.164443Z', 'text': 'SDXLモデルのDreambooth', 'unique_id': '5'}, {'create_date': '2024-03-05T14:07:56.264422Z', 'text': 'GPT×マーメイド記法で図やフローを出力', 'unique_id': '6'}, {'create_date': '2024-03-05T14:07:56.862481Z', 'text': 'WindowsでのDQN（Deep Q-learning Network）', 'unique_id': '12'}, {'create_date': '2024-03-05T14:07:56.566745Z', 'text': '動画からOpenpose用の骨格データをPNGで出力', 'unique_id': '9'}, {'create_date': '2024-03-05T14:07:56.373351Z', 'text': 'DETR（End-to-End Object Detection with Transformers）を使った動画での物 体検知', 'unique_id': '7'}, {'create_date': '2024-03-05T14:07:56.769679Z', 'text': 'DETR（End-to-End Object Detection with Transformers）物体検知', 'unique_id': '11'}, {'create_date': '2024-03-05T14:07:56.945689Z', 'text': 'Google ColabでのDQN（Deep Q-learning Network）', 'unique_id': '13'}, {'create_date': '2024-03-05T14:07:55.878655Z', 'text': 'sd-webui-animatediffで高解像 度・高フレームレートのクロマキーの動画を作成', 'unique_id': '2'}, {'create_date': '2024-03-05T14:07:55.779704Z', 'text': 'sd-webui-animatediffで高解像度・高フレームレートのクロマキーの動画 を作成（v3_sd15_mm.ckpt）', 'unique_id': '1'}, {'create_date': '2024-03-05T14:07:56.472649Z', 'text': 'DETR（End-to-End Object Detection with Transformers）物体検知（自前データでのファイン チューニング）', 'unique_id': '8'}, {'create_date': '2024-03-05T14:07:56.669887Z', 'text': 'DQN（Deep Q-learning Network）のハイパーパラメータ更新', 'unique_id': '10'}, {'create_date': '2024-03-05T14:07:55.975519Z', 'text': 'M1 MACで深層強化学習（スペック的には厳しい）', 'unique_id': '3'}]}}}

Weaviateの検索機能を使ったセマンティック検索を実行

from transformers import BertModel, BertJapaneseTokenizer
import torch
import weaviate
from datetime import datetime

# モデルとトークナイザーの初期化
model_name = "cl-tohoku/bert-base-japanese"
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)

# CUDAが利用可能なら使用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 単語「OpenCV」をエンベディング
text = "OpenCV"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
with torch.no_grad():
    outputs = model(**inputs)

embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()

# Weaviateクライアントの初期化
client = weaviate.Client("http://localhost:8080")

# エンベディングに最も近いテキストを検索するためのクエリを作成
# 厳しい条件でのクエリ
high_certainty_query = {
    "vector": embedding.tolist()[0],
    "certainty": 0.9  # 高い確信度
}

high_certainty_result = client.query.get("AddColumnTextData", ["text", "unique_id", "create_date"]).with_near_vector(high_certainty_query).with_additional(['certainty']).do()
print("High certainty result:", high_certainty_result)

以下の結果が得られる。
certaintyの値で類似度は確認できるが、直感的にここまで類似度が高いかと言われれば、微妙だと思うので、エンベディングモデル自体の性能は、高いとは言えない。

High certainty result: {'data': {'Get': {'AddColumnTextData': [{'_additional': {'certainty': 0.922598123550415}, 'create_date': '2024-03-05T14:07:56.164443Z', 'text': 'SDXLモデルのDreambooth', 'unique_id': '5'}, {'_additional': {'certainty': 0.9165770411491394}, 'create_date': '2024-03-05T14:07:56.862481Z', 'text': 'WindowsでのDQN（Deep Q-learning Network）', 'unique_id': '12'}, {'_additional': {'certainty': 0.9154360890388489}, 'create_date': '2024-03-05T14:07:56.945689Z', 'text': 'Google ColabでのDQN（Deep Q-learning Network）', 'unique_id': '13'}]}}}

感想

エンベディングモデルの精度の問題と思うが、類似度は納得のいく結果ではない。
実用的な精度を目指すには、この構成をたたき台とし、他のモデルを検証する必要がある。

その他

UUIDで取得したデータの削除

import weaviate

# Weaviateクライアントの初期化
client = weaviate.Client("http://localhost:8080")

# `AddColumnTextData`クラスのすべてのオブジェクトのUUIDを取得
query_result = client.query.get("AddColumnTextData", ["_additional { id }"]).do()

# オブジェクトが見つかった場合、それらを削除
if query_result["data"]["Get"]["AddColumnTextData"]:
    for obj in query_result["data"]["Get"]["AddColumnTextData"]:
        uuid = obj["_additional"]["id"]  # UUIDを取得
        # 各オブジェクトをUUIDを使用して削除
        client.data_object.delete(uuid, "AddColumnTextData")
        print(f"Deleted object with UUID: {uuid}")
else:
    print("No objects found to delete.")

クラスをドロップ

# WeaviateのPythonクライアントをインポートします。
import weaviate

# Weaviateサーバーへの接続を初期化します。ここではローカルホスト上の標準ポート(8080)を指定しています。
client = weaviate.Client("http://localhost:8080")

# 'TextData'クラスを削除します。
client.schema.delete_class("TextData")

クラス全件取得

import weaviate

client = weaviate.Client("http://localhost:8080")

query = """
{
  Get {
    TextData {
      text
      vector
    }
  }
}
"""

result = client.query.raw(query)
print(result)

お手製コサイン類似度検索

# 必要なライブラリをインポートします。
from transformers import BertModel, BertJapaneseTokenizer
import torch

# 使用するモデルの名称を指定します。ここでは、日本語対応のBERTモデル「cl-tohoku/bert-base-japanese」を使用します。
model_name = "cl-tohoku/bert-base-japanese"

# 指定したモデル名で、トークナイザー（テキストをトークンに分割するためのツール）を初期化します。
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)

# 指定したモデル名で、BERTモデルを初期化します。
model = BertModel.from_pretrained(model_name)

# エンベディングを生成したいテキストデータです。
text = "あなたのテキストをここに入れてください。"

# トークナイザーを使用して、テキストデータをモデルが理解できる形式に変換します。
# return_tensors="pt"は、PyTorchテンソルを出力するよう指定します。
# padding=Trueは、入力の長さを揃えるためにパディングを行うことを指定します。
# truncation=Trueは、指定された最大長さを超える入力を切り詰めることを指定します。
# max_length=128は、入力の最大長さを指定します。
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)

# モデルを使用して、入力からエンベディングを生成します。
outputs = model(**inputs)

# 最後の隠れ層の状態から平均を取り、テキストのエンベディングを取得します。
# .mean(dim=1)は、トークンごとのエンベディングの平均を計算します。
# .detach()は、計算グラフからこのテンソルを分離して、以降の計算で勾配が不要であることを示します。
# .numpy()は、テンソルをNumPy配列に変換します。
embeddings = outputs.last_hidden_state.mean(dim=1).detach().numpy()

# WeaviateのPythonクライアントをインポートします。
import weaviate

# Weaviateサーバーへの接続を初期化します。ここではローカルホスト上の標準ポート(8080)を指定しています。
client = weaviate.Client("http://localhost:8080")

# Weaviateのスキーマを定義します。このスキーマには、TextDataクラスが含まれています。
schema = {
    "classes": [{
        "class": "TextData", # クラス名を定義します。
        "vectorizer": "none",
        "properties": [{ # クラスのプロパティを定義します。
            "name": "text", # テキストデータを保持するプロパティです。
            "dataType": ["string"], # データタイプは文字列です。
            "description": "テキストデータ"
        },{
            "name": "vector",
            "dataType": ["number[]"],
            "description": "Vector representation of the text",
            "indexInverted": False  # 通常、ベクトルデータにはインバーテッドインデックスは不要です
        }]
    }]
}

# 定義したスキーマをWeaviateに作成します。
client.schema.create(schema)

# Weaviateにインサートするデータオブジェクトを定義します。
# `text`と`embedding`の値は、このコードの前に定義されている必要があります。
data_object = {
    "text": text, # テキストデータ。
    "vector": embeddings.tolist()[0] # エンベディングデータ。Numpy配列をリストに変換しています。
}

# 定義したデータオブジェクトをWeaviateのTextDataクラスにインサートします。
client.data_object.create(data_object, "TextData")


text2 = "あなたのテキストをここに入れます。"
inputs2 = tokenizer(text2, return_tensors="pt", padding=True, truncation=True, max_length=128)
outputs2 = model(**inputs2)
embeddings2 = outputs2.last_hidden_state.mean(dim=1).detach().numpy()

select_query = """
{
  Get {
    TextData {
      text
      vector
    }
  }
}
"""

# Weaviateのクエリ機能を使用して、特定のエンベディングに最も近いテキストデータを検索します。
# query_result = client.query.get("TextData", ["text"]).with_near_vector({"vector": embeddings2.tolist()[0]}).with_limit(2).with_additional(['certainty']).do()
# query_result = client.query.get("TextData", ["text"]).with_near_vector({"vector": client.query.raw(select_query)["data"]["Get"]["TextData"][0]["vector"]}).with_limit(2).with_additional(['certainty']).do()



from scipy.spatial.distance import cosine
db_data = client.query.raw(select_query)["data"]["Get"]["TextData"]
db_data_cosine_result = ""
db_data_cosine_result_value = 0
search_data = embeddings2.tolist()[0]
print(db_data)
for target_db_data in db_data:
    print("1")
    cosine_similarity = 1 - cosine(search_data, target_db_data["vector"])
    if cosine_similarity > db_data_cosine_result_value:
        db_data_cosine_result = target_db_data["text"]
        db_data_cosine_result_value = search_data

print(db_data_cosine_result)

# 検索結果を出力します。この結果には、クエリに最も近いテキストデータが含まれます。
# print(query_result)

import json
# print(json.dumps(query_result, indent=4))