RAGのベクトル検索のためにWeaviateを使ってみた

こんにちは、ぺいぺいです。検索を使用したLLMのRAGを実装しようとする際に、Vector DB と Vector Search が必要になったので、オープンソースのWeaviateを使用して、試して見ようという試みです。

基本知識

Weaviateとは

Weaviate は"The AI-native database for a new generation of software"と称されており、詰まるところ、AI活用を効率化するためのオープンソースベクトルデータベースとその周辺機能を提供するサービスである。ベクトルデータベースとは、テキストや画像などの情報をベクトルとして保存するデータベース（DB）であり、単語や画像の意味的な検索を可能にする。

RAGとは

Retrieval augmented generation (RAG) とは、他のドキュメントやテキストを参考にして、LLMの文章生成を行う手法であり、モデルがそれっぽい嘘をついてしまうHallusinationを低減させる方法として知られている。ChatGPTやGeminiなどのWebやPDFの参照機能がこれにあたる。Weaviateを使用すれば、RAGシステムを簡単に構築することができる。

NRIさんのサイトより引用（https://www.nri.com/jp/knowledge/glossary/lst/alphabet/rag）

RAGは検索（Retrieve）フェーズと生成（Generation）のフェーズに分けることができる。今回は、ReatrieveフェーズをWeaviateのベクトルDBとOpenAIのAPIの埋め込みモデルを使用して実現し、GenerationフェーズをOpenAI-APIのGPTシリーズを使用して実現する。

動作環境

• MacOS, Sonoma 14.6.1
• Apple M1 Max
• Python 3.12.3

環境構築

まずは、以下のサイトに従ってみる。

Weaviate Python client

まずは、以下コマンドを実行。Version 4.7.1 をインストール。

pip install -U weaviate-client

Pythonを実行し、Importしたら以下の警告が表示されたが一旦無視して進める。

.venv/lib/python3.12/site-packages/google/protobuf/runtime_version.py:112: UserWarning: Protobuf gencode version 5.27.2 is older than the runtime version 5.28.0 at grpc_health/v1/health.proto. Please avoid checked-in Protobuf gencode that can be obsolete.

Create a Weaviate instance

公式Webサイトには、Weaviate inctance か Docker instance かを選べと書いてあった。

• Weaviate Cloud（WCD） instance：環境構築やメンテナンスの管理を行いたくない人向け
• Docker instance：Weaviate をローカルマシンで動かしたい人向け

今回は、最終的にRAGを動かすことが目的であり、LLMの推論が必要になる。WCDでもAPIキーを使用したLLMの推論は可能のよう。なんとなくローカルで動かしてみたいので、久々にDockerを使う。よって、以下のサイトに従う。

Docker の環境構築

使用しているPCにはDocker環境が整っていなかったので、Docker の環境構築から始める。公式サイトからDocker Desktop for Mac をダウンロードする。

上記記事に従って、Dockerをinstall。
その後、作業ディレクトリに以下のdocker-compose.ymlを配置。

---
services:
  weaviate_anon:
    command:
    - --host
    - 0.0.0.0
    - --port
    - '8080'
    - --scheme
    - http
    image: cr.weaviate.io/semitechnologies/weaviate:1.26.1
    ports:
    - 8080:8080
    - 50051:50051
    restart: on-failure:0
    environment:
      OPENAI_APIKEY: $OPENAI_APIKEY
      QUERY_DEFAULTS_LIMIT: 25
      AUTHENTICATION_ANONYMOUS_ACCESS_ENABLED: 'true'
      PERSISTENCE_DATA_PATH: '/var/lib/weaviate'
      DEFAULT_VECTORIZER_MODULE: 'none'
      ENABLE_MODULES: 'text2vec-cohere,text2vec-huggingface,text2vec-openai,generative-openai,generative-cohere'
      BACKUP_FILESYSTEM_PATH: '/var/lib/weaviate/backups'
      CLUSTER_HOSTNAME: 'node1'
...

その後、作業ディレクトリでターミナルを開き、以下のコマンドで、Weacciate instance を作成。

docker compose up

この後、ドキュメントの指示に従って、http://localhost:8080にアクセスすると以下のJSONっぽいデータが出現。インスタンスを停止すると、アクセスできなくなるから多分成功している。

その後、自身のWeaviate instance のに接続できているかを確認するために、以下のpythonプログラムを実行し、接続が確認された。

import weaviate

client = weaviate.connect_to_local()
assert client.is_live()
print('connected')

Weaviate を用いたベクトル検索 + RAGの実装

Communicate with Weaviate

WeaviateとPythonを使用してどうやってやり取りするのかの基礎をさらう。上記の手順で、WeaviateのDockerコンテナを立てている状態からスタート。以下のPythonファイルを実行する。try-finally blockを使用することが推奨されている。

import weaviate
import json

try:
    client = weaviate.connect_to_local()
    assert client.is_live()
    print('------connected------')

    metainfo = client.get_meta()  # hostnameや埋め込みモデルなどのメタ情報
    print(json.dumps(metainfo, indent=2))
finally:
    # クライアントを閉じることでリソースを解放
    # 確実にcloseすることを保証
    client.close()

以下のようなメタ情報が出力される。

{
  "hostname": "http://[::]:8080",
  "modules": {
    "generative-cohere": {
      "documentationHref": "https://docs.cohere.com/reference/chat",
      "name": "Generative Search - Cohere"
    },
    "generative-openai": {
      "documentationHref": "https://platform.openai.com/docs/api-reference/completions",
      "name": "Generative Search - OpenAI"
    },
    "text2vec-cohere": {
      "documentationHref": "https://docs.cohere.ai/embedding-wiki/",
      "name": "Cohere Module"
    },
    "text2vec-huggingface": {
      "documentationHref": "https://huggingface.co/docs/api-inference/detailed_parameters#feature-extraction-task",
      "name": "Hugging Face Module"
    },
    "text2vec-openai": {
      "documentationHref": "https://platform.openai.com/docs/guides/embeddings/what-are-embeddings",
      "name": "OpenAI Module"
    }
  },
  "version": "1.26.1"
}

データベースを配置する

• 学ぶこと：collection を設定、作成した後、テキストデータをバッチinportで使用する
• できるようになること：collection を設定し、埋め込みモデルをセットできる。collection オブジェクトを作れるようになる。データをbatch import できるようになる

準備

このセクションでは、 Weaviate インスタンスをデータセットとともに配置し、OpenAIのAPIを使用して、テキストデータを埋め込めるようにする。今回必要なものは以下の3点。

• Weaviate instance：先述の方法で作成可能
• OpenAI key：チュートリアルに従うために必要。まだ作成したことがなかったので作成してみる
• ソースデータ：チュートリアルでは英語データセットを使用しているが、まったく同じなのはつまらないので、「農業基本オントロジー」データというオープンデータセットを使用してみる（CSV形式でダウンロード可能）。なお、有料APIを使用するので、使う行と列は減らしてデータサイズを小さくして使用する

まだOpenAIのAPIキーを作成したことがなかったのでこのタイミングで作成しておく（いずれ使うことになりそうなので）。OpenAIのサイトへ行き、APIを作成後、.zshrc に 環境変数OPENAI_API_KEYとして登録した。APIを使用するには、使用する前にお金を入金して残高を増やしておく必要があるので注意。以下のPythonスクリプトでGPTによる解答が返ってくれば動作していることが確認できる。

from openai import OpenAI
client = OpenAI()

completion = client.chat.completions.create(
  model="gpt-3.5-turbo",
  messages=[
    {"role": "system", "content": "あなたは優秀な日本語のアシスタントです。"},
    {"role": "user", "content": "日本の首都を漢字2文字で教えてください。"}
  ]
)

print(completion.choices[0].message)  # 「東京」関連の答えが返ってきた

Collection の作成

Weaviate は　"collections" というところにデータを保存する。collection は、同じデータ構造を共有するオブジェクトの集合である。ここでは、農作業オントロジー collectin を作成してみる。

import weaviate
import weaviate.classes.config as wc
import os

api_key = os.environ.get('OPENAI_API_KEY')
headers = {"X-OpenAI-Api-Key": api_key}  # Replace with your OpenAI API key
client = weaviate.connect_to_local(headers=headers)
print('------connected------')
client.collections.delete("Agriculture")

try:
    client.collections.create(
        name="Agriculture",
        properties=[
            wc.Property(name="task_name", data_type=wc.DataType.TEXT),
            wc.Property(name="purpose", data_type=wc.DataType.TEXT),
            wc.Property(name="action", data_type=wc.DataType.TEXT),
            wc.Property(name="target", data_type=wc.DataType.TEXT),
            wc.Property(name="location", data_type=wc.DataType.TEXT)
        ],
        # Define the vectorizer module
        vectorizer_config=wc.Configure.Vectorizer.text2vec_openai(),
        # Define the generative module
        generative_config=wc.Configure.Generative.openai()
    )
finally:
    client.close()

上記のプログラムを動かすことで、connectionが作成される。なお、複数回実行すると同名のconnectionは作成できないと怒られてエラーになるので注意。それぞれのconnectionには名前をちゃんと定義する必要がある。上記のコードの説明は以下の通り。

Properties

Properties はcollectionに保存したいオブジェクトの属性である。それぞれのPropertiesは名前とデータの型を持つ。今回のデータは全てテキストを想定しているので、TEXTとしているが、他にはINTやint型の配列を持つこともできる。

なお、Weaviateは自動の型推定も可能だが、明示的に示す方が良い。

Vectorizer configuration

もし埋め込みモデルを指定しない場合、Weaviateは適した埋め込みモデルを設定してくれる、今回はtext2vec-openiを指定する。OpenAPIのText-embedding-ada-002-v2が埋め込みモデルとして使用される。

Generative configuration

collectionはLLMなどのgenerative model と使用することができる。今回の例では、openaiモジュールを使用している。正式な名称はgenerative-openaiである。

Pythonのクラス

このコードの例では、Property, Datatype, Configureクラスを使用しているが、これらは全てweviate.classes.configで定義されている。

データのインポート

以下のプログラムを使用して、農作業オントロジーデータのcsvファイルをinportする。

import weaviate
import pandas as pd
from weaviate.util import generate_uuid5
from tqdm import tqdm
import os

uuid_list = []
api_key = os.environ.get('OPENAI_API_KEY')
headers = {"X-OpenAI-Api-Key": api_key}  # Replace with your OpenAI API key
client = weaviate.connect_to_local(headers=headers)

csv_path = './datasets/aao_mini_sample.csv'
df = pd.read_csv(csv_path)

# Get the collection
agris = client.collections.get("Agriculture")

try:
    # Enter context manager
    with agris.batch.dynamic() as batch:
        # Loop through the data
        for i, batch_df in tqdm(df.iterrows()):
            # Build the object payload
            agri_obj = {
                "task_name": batch_df["TASK_NAME"],
                "purpose": batch_df["PURPOSE_ja"],
                "action": batch_df["ACTION_ja"],
                "target": batch_df["TARGET_ja"],
                "location": batch_df["LOCATION_ja"],
            }

            # Add object to batch queue
            batch.add_object(
                properties=agri_obj,
                uuid=generate_uuid5(batch_df["TASK_NAME"]),
                # references=reference_obj  # You can add references here
            )
            
            # Note these are outside the `with` block - they are populated after the context manager exits
            failed_objs_a = client.batch.failed_objects  # Get failed objects from the first batch import
            failed_refs_a = client.batch.failed_references  # Get failed references from the first batch import

    # Check for failed objects
    if len(agris.batch.failed_objects) > 0:
        print(f"Failed to import {len(agris.batch.failed_objects)} objects")
    
finally:
    client.close()

Batch context マネージャー

batchオブジェクトはオブジェクトをbatcherに追加するためのものである。大規模データセットがある場合、バッチサイズの調整と時間調整を自動で行ってくれるので便利である。

.dynamic()メソッドは、データをimportしている最中にバッチサイズの決定と変更を自動で行ってくれるメソッドである。.fixed_size()メソッドはバッチサイズを固定化するものである。

データをBatcherに加える

データを加える前に、先ほどconnectionオブジェクトを作成する際に指定したデータの型にあらかじめ揃えておく必要が、今回のデータはすでに全てテキストデータで記載されているので、このプロセスは必要ない。

このようにして、ループによるバッチ処理で各オブジェクトをbacherに加えていく。batch.add_objectメソッドは、batcherにオブジェクトを追加するためのものである。そしてbatcherはバッチをWeaviate（collection）に送信する。

エラーハンドリング

バッチにはさまざまなオブジェクトが含まれ、いくつかのオブジェクトはimportに失敗してしまう可能性がある。どのオブジェクトでエラーが発生したのかをエラー処理で拾うことができる。

if len(movies.batch.failed_objects) > 0:
    print(f"Failed to import {len(movies.batch.failed_objects)} objects")

client.close()

ベクトル検索を行う

• 学ぶこと：Semantic Search（=意味検索）, Keyword Searcc（キーワード検索）、Hybrid Search（ハイブリッド検索）
• できるようになること：Semantic, Keyword, Hybrid Search の違いを理解し、それぞれを支えるようになる

Semantic search

Weaviateを使用すると、文章や単語の意味を用いて検索を行う semantic search を簡単に使用することができる。コードは以下の通り。

import weaviate
import os
import weaviate.classes.query as wq


api_key = os.environ.get('OPENAI_API_KEY')
headers = {"X-OpenAI-Api-Key": api_key}  # Replace with your OpenAI API key
try:
    client = weaviate.connect_to_local(headers=headers)

    agris = client.collections.get("Agriculture")

    response = agris.query.near_text(
        query='rice', limit=2, return_metadata=wq.MetadataQuery(distance=True)
    )
    print(response)

    for o in response.objects:
        print(
            o.properties['task_name']
        )
        print(
            f"Distance to query: {o.metadata.distance:.3f}\n"
        )  # Print the distance of the object from the query
finally:
    client.close()

上記のプログラムの結果は、クエリとデータベースのオブジェクトをともに埋め込んで作成した際のベクトルの類似度に基づいたものである。埋め込みモデルは先ほど記載したOpenAIのAPIを使用している。limitパラメーターは、返す検索結果の最大数を表し、return_metadataパラメータは、MetadataQueryクラスのオブジェクトを受け取るかを決めるものである。上記のプログラムでは、クエリと返ってきたオブジェクトとのベクトル間の距離を表している。

Semantic Search の検索対象

Semantic search では、与えられたテーブルデータの各列をどのように処理して埋め込んでいるのかを調べる。以下の公式ドキュメントを参照した。

Weaviateは、ベクトルの埋め込み表現を個々の"Properties"ごとではなく、オブジェクトごとに行なっている。例えば、プロパティ1が列1、プロパティ2が列2であるならば、列1、列2のテキストを合わせたテキストをまとめて埋め込む（各オブジェクト=各行）。プロパティ1がドキュメント、プロパティ2がドキュメントの要約であるならば、ドキュメントとその要約を合わせたテキストをまとめて埋め込む（各オブジェクト=ドキュメントと要約のペア）。特に特殊な設定をしていない場合の埋め込むためのテキストの結合規則は以下の通り。

• Textデータのみを使用して埋め込む
• プロパティ（=列)をテキストの結合前にアルファベット順に並べ替える
• vectorizePropertyName がtrueの時（デフォルトはfalse）プロパティ名（=列名）をそれぞれの列のテキストの前に追加する
• 各プロパティをスペースで結合
• クラス名を結合した文字列の先頭に追加（vectorizeClassNameがfalseでなければ）
• 生成された文字列を小文字に変換

例えば、以下のようなオブジェクト（行に該当）があるとする。

Article = {
  summary: "Cows lose their jobs as milk prices drop",
  text: "As his 100 diary cows lumbered over for their Monday..."
}

このオブジェクトは以下の文字列として埋め込まれる。

article cows lose their jobs as milk prices drop as his 100 diary cows lumbered over for their monday...

デフォルトでは、collection nameとプロパティのvaluesは計算に含められるが、プロパティのnamesは含まれない。つまり、テーブル全体の名前と各マスの文字列は含められるが、列名は含まれないということである。

プロパティごと（列ごと）に独立に埋め込みを行いたい場合は、skipとvectorizePropertyNameを使用しろと記載されているが、本記事ではここまでは深入りしない。

RAGを使用してみる

Retrieval augmented generation (RAG) とは、Semantic search の情報取得能力とLLMの文章生成能力を組み合わせた手法であり、モデルがそれっぽい嘘をついてしまうHallusinationを低減させる方法として知られている。Weaviateを使用すれば、RAGシステムを簡単に構築することができる。

設定

connection を作成する際に以下のパラメータを指定する。

generative_config=wc.Configure.Generative.openai()

ここでは、LLMモデルとして何を使用するかを設定することができるが、上記のパラメータを使用するとOpenAIのGPTシリーズが選択される。やはり、APIキーを取得しておく必要がある。

RAGクエリ

RAGクエリは、generativeクエリとWeaviateでは呼ばれており、各generativeクエリは、通常の検索クエリに加えて機能し、取得された各オブジェクトに対してRAGクエリを実行する。

シングルプロンプト生成

single prompt generation（シングルプロンプト生成）は、RAGを検索でヒットしてとってきたそれぞれのオブジェクトに対して適用する方法である。検索結果を比較せず、それぞれを個別に説明してもらいたい場合などに便利。プログラムの具体例は以下の通り。

import os
import weaviate
import os

api_key = os.environ.get('OPENAI_API_KEY')
headers = {"X-OpenAI-Api-Key": api_key}  # Replace with your OpenAI API key

client = weaviate.connect_to_local(headers=headers)

# Get the collection
agris = client.collections.get("Agriculture")

# Perform query
response = agris.generate.near_text(
    query="rice",
    limit=2,
    single_prompt="次の作業名の説明をしてください: {task_name}"
)

# Inspect the response
for o in response.objects:
    print(o.properties["task_name"])  # Print the title
    print(o.generated)  # Print the generated text (the title, in French)

client.close()

出力結果は以下のようになる。なお、生成にはデフォルトでは、GPT-3.5 turbo が使用されているようであり、上記プログラムのlimitを増やしすぎると、OpenAI API の時間あたりのリクエスト上限に引っかかってしまうので注意が必要。

鉄コーティング直播
鉄コーティング直播とは、鉄製品や鉄製部品に表面にコーティングを施す作業のことです。この作業は、鉄製品を錆や腐食から保護し、耐久性や耐候性を向上させるために行われます。鉄コーティング直播では、鉄製品を清掃し、表面処理を行った後、特殊な塗料やコーティング剤を塗布して鉄製品の表面を保護します。この作業は、建築や自動車産業などさまざまな産業で利用されています。
採種
「採種」とは、植物の種子を収穫する作業のことを指します。植物の種子を収穫し保存することで、将来的に新しい植物を育てることができます。採種は、植物の繁殖や保存、研究などさまざまな目的で行われます。採種の方法には、手作業での収穫や機械を使用した収穫などがあります。

グループタスク生成

Group task Generation（グループタスク生成）は、検索結果を全て同一のプロンプト生成に含めるタスクであり、検索結果を比較する場合や、全て一度に考慮して文章を生成したい場合などに便利。プログラムは以下の通り。

import os
import weaviate
import os

api_key = os.environ.get('OPENAI_API_KEY')
headers = {"X-OpenAI-Api-Key": api_key}  # Replace with your OpenAI API key

client = weaviate.connect_to_local(headers=headers)

# Get the collection
agris = client.collections.get("Agriculture")

query = '米'

try:
    # Perform query
    response = agris.generate.near_text(
        query=query,
        limit=2,
        grouped_task="次の作業名の中から最も米に関連する作業を選んで説明してください"  
    )

    # Inspect the response
    for o in response.objects:
        print(o.properties["task_name"])  # Print the title
    print(response.generated)  # Print the generated text (the title, in French)

finally:
    client.close()

出力例は以下の通り。

苗箱は種
種子繁殖作業
「苗箱は種」の作業名が最も米に関連する作業です。この作業は、米の種子を苗箱に植える作業を指します。種子を適切に植えることで、苗が育ち、最終的に収穫される米の生産が行われます。この作業は、米の生産過程において非常に重要な作業であり、品質の高い米を生産するために欠かせない作業です。

最後に

最後まで読んでいただきありがとうございました。何か間違い等がありましたらご指摘いただけますと幸いです。