初心者向け！LanceDBとLlamaIndexで始めるAI検索エンジン構築

Maki@Sunwood.ai.labs

2024年7月7日 18:29

近年のAI技術の進歩により、従来の検索エンジンを超える、より高度な検索体験を提供することが可能になりました。本記事では、その中核技術であるベクトル検索を、LanceDB と LlamaIndex を用いて実装する方法を、初心者の方でも理解できるように丁寧に解説します。

この記事で学べること

ベクトル検索の基本
LanceDBとLlamaIndexの概要
LanceDBとLlamaIndexを用いたベクトル検索の実装
Pythonコードと詳細なコメントによる解説

参考資料

ベクトル検索とは？

ベクトル検索は、テキストや画像などのデータを多次元ベクトルに変換し、そのベクトル間の類似度に基づいて検索を行う技術です。従来のキーワード検索では不可能だった、意味的な類似性に基づいた検索を可能にします。

例えば、「美味しいラーメン屋」というクエリに対して、キーワード検索では「美味しい」「ラーメン屋」といった単語が完全に一致するページしかヒットしません。一方、ベクトル検索では、「ラーメン」「つけ麺」「まぜそば」といった、意味的に関連性の高い単語を含むページも検索結果に表示することができます。

LanceDBとLlamaIndex：ベクトル検索を容易にする強力なツール

LanceDB は、ベクトルデータを高速に検索するためのオープンソースのベクトルデータベースです。大規模なデータセットにも対応しており、検索の精度と速度を両立しています。

LlamaIndex は、大規模言語モデルを活用したアプリケーション開発を容易にするためのフレームワークです。テキストデータの処理、ベクトル化、様々なベクトルデータベースとの連携など、ベクトル検索に必要な機能を提供しています。

LanceDBとLlamaIndexを用いたベクトル検索の実装

それでは、実際にLanceDBとLlamaIndexを用いてベクトル検索を実装してみましょう。

環境設定

まず、必要なライブラリをインストールします。

!pip install llama-index llama-index-vector-stores-lancedb lancedb openai torch transformers

!pip install tantivy

データの準備

今回は、Paul Grahamのエッセイをサンプルデータとして使用します。以下のコードでデータをダウンロードし、テキストファイルとして保存します。

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

インデックスの作成

LlamaIndexを使って、ダウンロードしたテキストデータを読み込み、LanceDBにインデックスを作成します。

import logging
import sys
import openai
import os

# OpenAI APIキーを設定
# openai.api_key = "YOUR_OPENAI_API_KEY" # ここにあなたのOpenAI APIキーを入力してください。
from google.colab import userdata

openai.api_key = userdata.get('OPENAI_API_KEY')

# ロギング設定
logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index.core import SimpleDirectoryReader, Document, StorageContext
from llama_index.core import VectorStoreIndex
from llama_index.vector_stores.lancedb import LanceDBVectorStore
import textwrap

# データの読み込み
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

# LanceDBベクトルストアの作成
vector_store = LanceDBVectorStore(
    uri="./lancedb", # データベースの保存先
    mode="overwrite", # 既存のデータベースを上書き
    query_type="hybrid" # ハイブリッド検索を有効化
)

# ストレージコンテキストの作成
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# インデックスの作成
index = VectorStoreIndex.from_documents(
    documents, storage_context=storage_context
)

検索の実行

作成したインデックスを使って、検索を実行してみましょう。

from llama_index.core.vector_stores import (
    MetadataFilters,
    FilterOperator,
    FilterCondition,
    MetadataFilter,
)

from datetime import datetime

# メタデータフィルタの作成
query_filters = MetadataFilters(
    filters=[
        MetadataFilter(
            key="creation_date",
            operator=FilterOperator.EQ,
            value=datetime.now().strftime("%Y-%m-%d"),
        ),
        MetadataFilter(
            key="file_size",
            value=75040,
            operator=FilterOperator.GT
        ),
    ],
    condition=FilterCondition.AND,
)

# 検索の実行
query_engine = index.as_query_engine(filters=query_filters)
response = query_engine.query("How much did Viaweb charge per month?")

# 結果の表示
print(response)
print("metadata -", response.metadata)

ハイブリッド検索と再ランク付け

LanceDBは、再ランク付け機能を備えたハイブリッド検索を提供します。完全なドキュメントについては、こちらを参照してください。

この例では、「colbert」再ランクを使用します。次のセルは、「colbert」に必要な依存関係をインストールします。別の再ランクを選択する場合は、それに応じて依存関係を調整してください。

from lancedb.rerankers import ColbertReranker

# 再ランクの作成
reranker = ColbertReranker()

# ベクトルストアへの再ランクの追加
vector_store._add_reranker(reranker)

# 検索エンジンの作成
query_engine = index.as_query_engine(
    filters=query_filters,
)

# 検索の実行
response = query_engine.query("How much did Viaweb charge per month?")

# 結果の表示
print(response)
print("metadata -", response.metadata)

データの追加

既存のインデックスにデータを追加することもできます。

from llama_index.core import VectorStoreIndex
from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo

# 既存のquery_engineを使用して検索を実行
response = query_engine.query("Where is the sky purple?")

# 結果の表示
print("Initial query response:")
print(response)
print("metadata -", response.metadata)

# テキストノードの作成
node1 = TextNode(text="The sky is purple in Portland, Maine", id_="portland_sky")
node2 = TextNode(text=str(response), id_="query_response")

# ノード間の関係を設定
node1.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(node_id=node2.node_id)
node2.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(node_id=node1.node_id)

# ノードのリストを作成
nodes = [node1, node2]

# 既存のインデックスを削除（必要な場合）
try:
    del index
except NameError:
    pass  # インデックスが存在しない場合は何もしない

# 新しいインデックスの作成
index = VectorStoreIndex(nodes)

# 新しい検索の実行
new_query_engine = index.as_query_engine()
new_response = new_query_engine.query("Where is the sky purple?")

# 新しい結果の表示
import textwrap
print("\nNew query response:")
print(textwrap.fill(str(new_response), 100))

# ノードの内容と関係を表示（デバッグ用）
print("\nNode contents and relationships:")
for node in nodes:
    print(f"Node ID: {node.node_id}")
    print(f"Text: {node.text}")
    print("Relationships:")
    for rel_type, rel_info in node.relationships.items():
        print(f"  {rel_type}: {rel_info.node_id}")
    print()

既存のテーブルからのインデックスの作成

既存のテーブルからインデックスを作成することもできます。

# 既存のインデックスを削除
del index

# 既存のテーブルからベクトルストアを作成
vec_store = LanceDBVectorStore.from_table(vector_store._table)

# ベクトルストアからインデックスを作成
index = VectorStoreIndex.from_vector_store(vec_store)

# 検索の実行
query_engine = index.as_query_engine()
response = query_engine.query("What companies did the author start?")

# 結果の表示
print(textwrap.fill(str(response), 100))

📒ノートブック

https://colab.research.google.com/drive/1VB1EaqoCAgTqhC5VZHM2HzfGW8c_Ik5I?usp=sharing

まとめ

本記事では、LanceDBとLlamaIndexを用いてベクトル検索を実装する方法を解説しました。これらのツールを使用することで、意味的な類似性に基づいた高度な検索システムを容易に構築することができます。

今後の展望

ベクトル検索は、検索エンジンだけでなく、チャットボット、レコメンドシステムなど、様々なアプリケーションで活用されています。今後も、AI技術の進化とともに、その応用範囲はますます広がっていくでしょう。

この記事が気に入ったらサポートをしてみませんか？