オープンソースのベクトルデータベースである Weaviate を Docker で起動しデータを投入し、そのデータを使って RAG (検索拡張生成) を行うチュートリアルです。

OpenAI の API を使うので、OpenAI の API キーが必要です。このチュートリアルをするにあたり、OpenAIの従量課金の使用料が発生します。

この記事は、 2024-10-11 に、TORICO の社内勉強会で行った内容の共有となります。

ベクトルDBとは

データを N次元ベクトルとして記録し、類似度検索ができるものです。

オープンソースで手軽に使えるのは、Weaviate や Chroma , Qdrant, Elasticsearch など。MySQL + ベクトル検索機能を各種ベンダーが提供していたりもします。

テキスト以外にも、画像や音声や動画などもベクトルデータにすれば類似度検索ができます。

Weaviate とは

https://weaviate.io/ja

オープンソースのベクトルデータベース。

テキストのベクトル化(Embedding) やRAGの生成モジュールに様々なモジュールをプラグインできる特徴があります。

例えば、OpenAI の API を使ったり、ローカルで動作している Transformer と連携したりといったことが柔軟に行えます。

Doker イメージで手軽に無料で動かせるので、社内 Kubernetes などでサービスを作ってコストを抑えたい場合に良いと思います。

プロジェクトフォルダの作成

開発用フォルダの中に、weaviate-tutorial フォルダを作ってください。

その中に、 .env ファイルを作り、OPENAI_APIKEY 環境変数を定義してください。

mkdir weaviate-tutorial
cd weaviate-tutorial
vim .env

フォルダ構成

weaviate-tutorial
 + .env

.env

OPENAI_APIKEY=sk-xxxxxxxxxxxxxxxxxxxxxxxx

環境構築

Pipfile

[[source]]
name = "pypi"
url = "https://pypi.org/simple"
verify_ssl = true

[dev-packages]

[packages]
weaviate-client = "*"
python-dotenv = "*"
requests = "*"

[requires]
python_version = "3.12"

pipenv install
pipenv shell

ベクトル化を体感する

まずは「テキストをベクトルにする」というイメージを得るために、OpenAI の エンベッティングAPI を使ってテキストをベクトルにしてみます。

mkdir scripts
vim scripts/search_utils.py

フォルダ構成

weaviate-tutorial
 + .env
 + Pipfile
 + scripts
    + search_utils.py

scripts/search_utils.py

import requests
import os
import dotenv

dotenv.load_dotenv()

OPENAI_APIKEY = os.environ.get('OPENAI_APIKEY', '')

if not OPENAI_APIKEY:
    raise RuntimeError('env OPENAI_APIKEY is required.')


def get_vector_by_openai(text) -> list[float]:
    """
    OpenAI の embedding API を使って文字列をベクトルに変換する
    https://platform.openai.com/docs/guides/embeddings
    """
    response = requests.post(
        'https://api.openai.com/v1/embeddings',
        headers={
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {OPENAI_APIKEY}'
        },
        json={
            'input': text,
            'model': 'text-embedding-3-small'
        }
    )

    response.raise_for_status()
    return response.json()['data'][0]['embedding']



if __name__ == '__main__':
    print(get_vector_by_openai('少年コミック誌の黄金期で一番おもしろい漫画'))

実行する

python3 scripts/search_utils.py

実行すると、テキストが 1,536次元のベクトル ( float のリスト )に変換されます。

近年のAI は、何かしらのデータをベクトルに変換して処理するのが大半です。 例えば GPT-3 では、単語を複数に分割した「トークン」1つを、12,288次元のベクトルとして扱います。

何かしらの入力を、学習されたパラメーターによって高次元のベクトルに変換することで、類似の検索や推論などの計算処理ができるようになります。

Docker 起動スクリプトの作成

今回は、docker-compose や kubernetes 等は使わず、 docker コマンド一発で起動します。

Text2Vector モジュールは、OpenAI の API を使います。

起動時に APIトークンを環境変数として指定すれば、Docker コンテナの中で OpenAI API を使ってくれます。

docker フォルダを作り、その中に起動用のシェルスクリプトを作ります。

データフォルダとしてマウントするための data フォルダを、docker フォルダの中に作ります。

mkdir docker
cd docker
mkdir data
vim run.sh

run.sh

#!/usr/bin/env zsh

cd $(dirname $0)

source ../.env

docker run \
  --rm \
  -p 8080:8080 \
  -p 50051:50051 \
  -v $(pwd)/data:/data \
  --env ENABLE_MODULES=text2vec-openai,generative-openai \
  --env OPENAI_APIKEY=${OPENAI_APIKEY} \
  cr.weaviate.io/semitechnologies/weaviate:1.26.4

試しに起動します。

chmod +x run.sh
./run.sh

フォルダ構成

 + .env
 + Pipfile
 + scripts
 |  + search_utils.py
 + docker
    + data
    + run.sh

(参考例) Text To Vecotr モジュール (Transformer) を自前で用意する場合

Text2Vector などのモジュールに、OpenAI のAPIを使わずに、自前で Transformer などを動かして使うこともできます。

その場合は複数の Docker コンテナを起動することになるので、docker コマンド一発ではなく docker compose を使うと良いでしょう。

docker-compose.yaml のサンプルが weaviate のページにあります。

https://weaviate.io/developers/weaviate/installation/docker-compose#weaviate-with-the-text2vec-transformers-module

サンプルデータを準備する

社内のマンガのデータベース(MySQL)から作成しました。

本ブログにはデータは添付していません。

サンプルとして生成用の SQL を添付します。MySQL の場合、JSON_* 関数を使うことで、リレーショナルデータのエクスポートが便利に行えます。

SELECT
    p.product_id,
    p.title,
    p.description,
    p.sku,
    p.price_tax_excluded,
    -- 作者名を JSON 配列として出力
    (SELECT JSON_ARRAYAGG(a.title)
     FROM product_author AS pa
     JOIN author AS a
     ON pa.author_id = a.id
     WHERE pa.product_id = p.product_id) AS authors,
    -- 出版社名を JSON 配列として出力
    (SELECT JSON_ARRAYAGG(p.title)
     FROM product_publisher AS pp
     JOIN publisher AS p
     ON pp.publisher_id = p.id
     WHERE dpc.product_id = dp.product_id) AS comics,
    -- タグを JSON 配列として出力
    (SELECT JSON_ARRAYAGG(t.title)
     FROM product_tag AS pt
     JOIN tag AS t
     ON pt.tag_id = t.id
     WHERE pt.product_id = p.product_id) AS tags,
    -- レビューを JSON オブジェクトの配列として出力 
    (SELECT JSON_ARRAYAGG(JSON_OBJECT(
        'title', pr.title,
        'score', pr.score,
        'comment', pr.comment))
     FROM product_review AS pr
     WHERE pr.product_id = p.product_id
       AND pr.active = 1) AS reviews
FROM product AS p
WHERE p.active = 1
AND p.description != ''
ORDER BY product_id
LIMIT 3000

Weaviate にコレクションを作る

コレクションというのは、RDBMS でいうテーブル、Elasticsearchでいう「インデックス」にあたるものです。

スクリプトファイルを作って実行します。

scripts/c01_define_data_collection.py

import weaviate
import weaviate.classes.config as wc
import weaviate.classes as wvc

client = weaviate.connect_to_local(port=8080)

try:
    questions = client.collections.create(
        name='Comic',
        # properties は指定しなくても動作するが、指定したほうが誤動作が無くて良い。
        properties=[
            wc.Property(name='product_id', data_type=wc.DataType.INT),
            wc.Property(name='title', data_type=wc.DataType.TEXT),
            wc.Property(name='description', data_type=wc.DataType.TEXT),
            wc.Property(name='sku', data_type=wc.DataType.TEXT),
            wc.Property(name='price', data_type=wc.DataType.INT),
            wc.Property(name='authors', data_type=wc.DataType.TEXT_ARRAY),
            wc.Property(name='publishers', data_type=wc.DataType.TEXT_ARRAY),
            wc.Property(name='tags', data_type=wc.DataType.TEXT_ARRAY),
            wc.Property(
                name='reviews', data_type=wc.DataType.OBJECT_ARRAY,
                nested_properties=[
                    wc.Property(name='title', data_type=wc.DataType.TEXT),
                    wc.Property(name='comment', data_type=wc.DataType.TEXT),
                    wc.Property(name='score', data_type=wc.DataType.INT),
            ]),
        ],
        # https://platform.openai.com/docs/models/embeddings
        # ada babbage curie davinci 〜第2世代
        # text-embedding-3-small 1,536次元
        # text-embedding-3-large 3,072次元 非英語に最適 smallの7倍近い金額
        vectorizer_config=wvc.config.Configure.Vectorizer.text2vec_openai(model='text-embedding-3-small'),
        # generative_config: gpt-4o は設定できるけど使えない。
        # gpt-4o-mini はそもそも設定できない。
        # 安全のためデフォルトで使う。
        # 後から設定変更できない。
        # available model names are:
        # gpt-3.5-turbo gpt-3.5-turbo-16k gpt-3.5-turbo-1106 gpt-4 gpt-4-32k gpt-4-1106-preview gpt-4o
        generative_config=wvc.config.Configure.Generative.openai()
    )
finally:
    client.close()

T2V の埋め込みモデルには OpenAI の text-embedding-3-small モデルを使っています。

英語以外の場合は、 text-embedding-3-large を使った方が良いとは思いますが、今回のチュートリアルではコストを抑えるため small にしています。

それより古い、第2世代以前のモデルは採用する理由は無いでしょう。

今後 RAG をするため、generative_config を指定しています。

wvc.config.Configure.Generative.openai は、引数でモデルを指定できます。gpt-4o を使いたいところですが、model='gpt-4o' を指定した場合コレクションは作れますが実際の生成時にエラーが出ます。

今回のチュートリアルでは、gpt-4 はコストと生成時間が許容できず、gpt-4o-mini は指定するとコレクションの生成時に失敗する(2024-09-29 現在)ため使えませんでした。

指定しないと gpt-3.5-turbo が使われます。今回は、 gpt-3.5 の品質に不満はありつつも、コストと今回のチュートリアル実行のスムーズな進行のためにモデルの指定を行わず、デフォルトの gpt-3.5-turbo を使うことにしました。

あと、執筆時点の Weaviate だと、generative_config は collections.create した後には変更できないようです。(埋め込みモデルなどは変更できるが生成モジュールを変更するAPIコードが無い。生成時点での指定もできない)

この問題においての Issue があり、PR はマージされていますが Issue は執筆時点では Open のままでした。

スクリプトを作成後、実行します。

データをロードする

ロード用のスクリプトを作って実行します。

scripts/c02_load_data.py

"""
データを投入する。

3000行で30秒ぐらいかかる。
"""
import time

import weaviate
import json

client = weaviate.connect_to_local(port=8080)

try:

    with open('./comics.json') as fp:
        table = json.load(fp)

    comic_collection = client.collections.get('Comic')
    start_time = time.time()
    print('c02_load_data inserting...')
    comic_collection.data.insert_many(table)
    elapsed_time = time.time() - start_time
    print(f'c02_load_data completed. elapsed {elapsed_time:.0f} seconds')

finally:
    client.close()

セマンティック検索

シンプルに検索してみます。

scripts/c03_semantic_search.py

"""
Semantic search
https://weaviate.io/developers/academy/py/starter_text_data/text_searches/semantic
"""
import weaviate
import weaviate.classes.query as wq

client = weaviate.connect_to_local(port=8080)

try:
    col = client.collections.get('Comic')

    response = col.query.near_text(
        query='熱血野球',
        limit=10,
        return_metadata = wq.MetadataQuery(distance=True)
    )

    for r in response.objects:
        print(r.properties, r.metadata.distance)
finally:
    client.close()

query 引数には、複数のキーワードを入れることができます。

文字列でも文字列のリストでも入れることができて、効果が少し違います。

文字列のリスト (['テニス', 'ブラックホール']) を入れると、どちらも等価なワードとして扱われますが、スペース区切りの単語 (テニス ブラックホール)の場合は語順を評価されます。

ベクトルDBなので、検索ワードをベクトル化して、コサイン類似度の近傍検索などしてるのだと思います。

キーワードサーチ

scripts/c04_keyword_search.py

"""
Keyword search
https://weaviate.io/developers/academy/py/starter_text_data/text_searches/keyword_hybrid
"""
from re import search

import weaviate
import weaviate.classes.query as wq

client = weaviate.connect_to_local(port=8080)

try:
    col = client.collections.get('Comic')

    search_method = col.query.bm25  # 純粋なキーワード検索
    # search_method = col.query.hybrid  # キーワードとセマンティック検索のハイブリッド
    response = search_method(
        query='スタンド',
        limit=10,
        return_metadata = wq.MetadataQuery(score=True)
    )

    for r in response.objects:
        print(r.properties, r.metadata.distance)
finally:
    client.close()

bm25 メソッドで単純なキーワード検索ができて、 hybrid だとキーワードとセマンティック検索のハイブリッドでの検索が行えます。

bm25 だけを使うのであれば、あまりベクトルDBを使う利点が無さそうに思います。

ハイブリッド検索は有用そうな気がします。

ベクトルデータを事前に作って検索

何らかの方法でベクトルデータ ( list[float] ) を事前に作り、それで直接検索することもできます。

今回は、 OpenAI の embedding API を使ってベクトルデータを作ってそれで検索してみます。

scripts/search_utils.py

(再掲)

import requests
import os
import dotenv

dotenv.load_dotenv()

OPENAI_APIKEY = os.environ.get('OPENAI_APIKEY', '')

if not OPENAI_APIKEY:
    raise RuntimeError('env OPENAI_APIKEY is required.')


def get_vector_by_openai(text) -> list[float]:
    """
    OpenAI の embedding API を使って文字列をベクトルに変換する
    """
    response = requests.post(
        'https://api.openai.com/v1/embeddings',
        headers={
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {OPENAI_APIKEY}'
        },
        json={
            'input': text,
            'model': 'text-embedding-3-small'
        }
    )

    response.raise_for_status()
    return response.json()['data'][0]['embedding']

scripts/c05_vector_search.py

"""
Keyword search
https://weaviate.io/developers/academy/py/starter_text_data/text_searches/keyword_hybrid
"""
import weaviate
import weaviate.classes.query as wq
from search_utils import get_vector_by_openai

client = weaviate.connect_to_local(port=8080)
vector = get_vector_by_openai('バレーボール インターハイ')

try:
    col = client.collections.get('Comic')
    response = col.query.near_vector(
        vector,
        limit=10,
        return_metadata = wq.MetadataQuery(score=True)
    )

    for r in response.objects:
        print(r.properties, r.metadata.distance)
finally:
    client.close()

RAG シングルプロンプト

検索結果の1レコードごとにAI生成を行います。

scripts/c11_rag_single_prompt.py

import weaviate
client = weaviate.connect_to_local(port=8080)

try:
    col = client.collections.get("Comic")

    question = 'スポーツ作品で双子が出てくるものはどんなのがある?'

    response = col.generate.near_text(
        query=question,  # 後述
        limit=10,
        single_prompt="""次のコミックの検索該当結果の見どころを要約してください。
検索ワード: %s

該当タイトル: {title}

作者: {','.join(authors) if authors else ''}

出版社: {','.join(publishers) if publishers else ''}

タグ: {','.join(tags) if tags else ''}

## 作品詳細
{description}

## レビュー
{reviews or '(レビュー無し)'}
""" % question
    )
    for r in response.objects:
        print('##', r.properties['name'])
        print(r.generated)
        print()

finally:
    client.close()

生成結果の1レコードごとに、single_prompt を元に生成が行われます。

single_prompt は python の文字列で、f-string 記法が適用されるので、簡単な文字列操作を書くことができます。

検索ワードは、自然文の質問を query に入れてますが、これで十分な精度が出るかはわかりませんでした。

一応、検索結果としてはそれっぽく該当するのですが、常に最適なものが出せているかは不明です。

自然文を query にそのまま入れる前に、LLMで「ベクトル検索用のキーワードを作って」といったプロンプトで変換して検索もしてみましたが、結果としては良し悪しだと思いました。

ベクトルの一致度としては、質問の文章「双子が登場するのは？」よりも「双子が登場する」といった断定形式に変更してクエリとして使ったほうが適切かもしれません。ここは経験が無く、精度を詰めきれませんでした。

RAG グループタスク

生成結果全体に対してRAG を行います。

scripts/c12_rag_grouped_task.py

import weaviate
client = weaviate.connect_to_local(port=8080)

try:
    col = client.collections.get("Comic")

    question = 'ハンター漫画の主人公の名前を教えて'

    response = col.generate.near_text(
        query=question,
        limit=5,
        grouped_task=question
    )
    for r in response.objects:
        print('##', r.properties['name'])
    print('-----------------------------------------------')
    print(response.generated)

finally:
    client.close()

grouped_task に、ユーザー入力の自然文を入れればそれっぽい返答になります。

検索文言のキーワードをLLMで作成する

query に直接自然文を入れるより、一回 LLM などで検索用のワードに変換したほうが精度が良くなるかもしれません。

私がやった感じでは、どっちもどっち…といった結果でした。

検索する元データの充実度が不十分だったからだと思います。

script/search_utils.py

import requests
import os
import dotenv

dotenv.load_dotenv()

OPENAI_APIKEY = os.environ.get('OPENAI_APIKEY', '')

if not OPENAI_APIKEY:
    raise RuntimeError('env OPENAI_APIKEY is required.')


def text_parse_to_search_words(text: str) -> list[str]:
    """
    テキストを検索ワードに変換する
    """
    response = requests.post(
        'https://api.openai.com/v1/chat/completions',
        headers={
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {OPENAI_APIKEY}'
        },
        json={
            'model': 'gpt-4o-mini',
            'messages': [
                {
                    'role': 'system',
                    'content': 'ユーザーから入力される検索用の質問文を解析して、'
                               'ベクトルデータベース(Weaviate) の検索ワードとして使う単語のリストに変換し、'
                               'カンマ区切りで返答しなさい'
                },
                {
                    'role': 'user',
                    'content': text
                }
            ],
        }
    )
    response.raise_for_status()
    content = response.json()['choices'][0]['message']['content']
    return list(filter(None, map(lambda w:w.strip(), content.split(','))))

script/search_utils.py

import weaviate
from search_utils import text_parse_to_search_words

client = weaviate.connect_to_local(port=8080)

try:
    col = client.collections.get("Comic")

    question = '大人でも楽しめるギャグ漫画を教えて'

    query_words = text_parse_to_search_words(question)
    print(query_words)
    response = col.generate.near_text(
        query=query_words,
        limit=10,
        grouped_task=question,
    )
    for r in response.objects:
        print('##', r.properties['name'])
    print('-----------------------------------------------')
    print(response.generated)


finally:
    client.close()