PythonでTiDB + AIを使い始める

このチュートリアルでは、セマンティック検索機能を備えたシンプルなAIアプリケーションの開発方法を説明します。従来のキーワード検索とは異なり、セマンティック検索はクエリの背後にある意味をインテリジェントに理解し、最も関連性の高い結果を返します。例えば、「犬」、「魚」、「木」というタイトルのドキュメントがあり、「泳ぐ動物」を検索すると、アプリケーションは「魚」を最も関連性の高い結果として特定します。

このチュートリアルでは、 TiDBベクトル検索 、Python、 Python 用 TiDB ベクター SDK 、および AI モデルを使用してこの AI アプリケーションを開発します。

注記：

ベクター検索機能はベータ版です。予告なく変更される可能性があります。バグを見つけた場合は、GitHubで問題報告を行ってください。

注記：

ベクトル検索機能は、TiDB Self-Managed、 TiDB Cloudサーバーレス TiDB Cloud専用利用できます。TiDB Self-ManagedおよびTiDB Cloud Dedicatedの場合、TiDBバージョンはv8.4.0以降である必要があります（v8.5.0以降を推奨）。

前提条件

このチュートリアルを完了するには、次のものが必要です。

• Python 3.8以上個インストールされました。
• ギット個インストールされました。
• TiDB クラスター。

TiDB クラスターがない場合は、次のように作成できます。

• (推奨) TiDB Cloud Serverless クラスターの作成に従って、独自のTiDB Cloudクラスターを作成します。
• ローカルテストTiDBクラスタをデプロイまたは本番のTiDBクラスタをデプロイに従って、v8.4.0 以降のバージョンのローカル クラスターを作成します。

始めましょう

以下の手順は、アプリケーションをゼロから開発する方法を示しています。デモを直接実行するには、 pingcap/tidb-vector-pythonリポジトリのサンプルコードをチェックアウトしてください。

ステップ1. 新しいPythonプロジェクトを作成する

任意のディレクトリに、新しい Python プロジェクトとexample.py名前のファイルを作成します。

mkdir python-client-quickstart
cd python-client-quickstart
touch example.py

ステップ2. 必要な依存関係をインストールする

プロジェクト ディレクトリで、次のコマンドを実行して必要なパッケージをインストールします。

pip install sqlalchemy pymysql sentence-transformers tidb-vector python-dotenv

• tidb-vector : TiDB ベクトル検索と対話するための Python クライアント。
• sentence-transformers : テキストからベクトル埋め込み生成するための事前トレーニング済みモデルを提供する Python ライブラリ。

ステップ3. TiDBクラスターへの接続文字列を構成する

選択した TiDB デプロイメント オプションに応じて、クラスター接続文字列を構成します。

• TiDB Cloud Serverless
• TiDB Self-Managed

TIDB_DATABASE_URL="mysql+pymysql://<USER>:<PASSWORD>@<HOST>:<PORT>/<DATABASE>"
# For example: TIDB_DATABASE_URL="mysql+pymysql://root@127.0.0.1:4000/test"

ステップ4. 埋め込みモデルの初期化

埋め込みモデルデータをベクトル埋め込みに変換します。この例では、テキスト埋め込みに事前学習済みのモデルmsmarco-MiniLM-L12-cos-v5使用します。7 sentence-transformersが提供するこの軽量モデルは、テキストデータを 384 次元のベクトル埋め込みに変換します。

モデルをセットアップするには、次のコードをexample.pyファイルにコピーします。このコードはSentenceTransformerインスタンスを初期化し、後で使用するtext_to_embedding()関数を定義します。

from sentence_transformers import SentenceTransformer

print("Downloading and loading the embedding model...")
embed_model = SentenceTransformer("sentence-transformers/msmarco-MiniLM-L12-cos-v5", trust_remote_code=True)
embed_model_dims = embed_model.get_sentence_embedding_dimension()

def text_to_embedding(text):
    """Generates vector embeddings for the given text."""
    embedding = embed_model.encode(text)
    return embedding.tolist()

ステップ5. TiDBクラスターに接続する

TiDBVectorClientクラスを使用して TiDB クラスターに接続し、ベクター列を持つテーブルembedded_documentsを作成します。

注記

表内のベクトル列の次元が、埋め込みモデルによって生成されるベクトルの次元と一致していることを確認してください。例えば、 msmarco-MiniLM-L12-cos-v5モデルは384次元のベクトルを生成するため、 embedded_documentsのベクトル列の次元も384である必要があります。

import os
from tidb_vector.integrations import TiDBVectorClient
from dotenv import load_dotenv

# Load the connection string from the .env file
load_dotenv()

vector_store = TiDBVectorClient(
   # The 'embedded_documents' table will store the vector data.
   table_name='embedded_documents',
   # The connection string to the TiDB cluster.
   connection_string=os.environ.get('TIDB_DATABASE_URL'),
   # The dimension of the vector generated by the embedding model.
   vector_dimension=embed_model_dims,
   # Recreate the table if it already exists.
   drop_existing_table=True,
)

ステップ6. テキストデータを埋め込み、ベクトルを保存する

このステップでは、「dog」、「fish」、「tree」といった単語を含むサンプルドキュメントを準備します。以下のコードは、 text_to_embedding()関数を使用してこれらのテキストドキュメントをベクトル埋め込みに変換し、ベクトルストアに挿入します。

documents = [
    {
        "id": "f8e7dee2-63b6-42f1-8b60-2d46710c1971",
        "text": "dog",
        "embedding": text_to_embedding("dog"),
        "metadata": {"category": "animal"},
    },
    {
        "id": "8dde1fbc-2522-4ca2-aedf-5dcb2966d1c6",
        "text": "fish",
        "embedding": text_to_embedding("fish"),
        "metadata": {"category": "animal"},
    },
    {
        "id": "e4991349-d00b-485c-a481-f61695f2b5ae",
        "text": "tree",
        "embedding": text_to_embedding("tree"),
        "metadata": {"category": "plant"},
    },
]

vector_store.insert(
    ids=[doc["id"] for doc in documents],
    texts=[doc["text"] for doc in documents],
    embeddings=[doc["embedding"] for doc in documents],
    metadatas=[doc["metadata"] for doc in documents],
)

ステップ7. セマンティック検索を実行する

このステップでは、既存のドキュメント内のどの単語にも直接一致しない「a swimming animal」を検索します。

次のコードでは、 text_to_embedding()関数を再度使用してクエリ テキストをベクター埋め込みに変換し、その埋め込みを使用してクエリを実行して、最も近い上位 3 つの一致を検索します。

def print_result(query, result):
   print(f"Search result (\"{query}\"):")
   for r in result:
      print(f"- text: \"{r.document}\", distance: {r.distance}")

query = "a swimming animal"
query_embedding = text_to_embedding(query)
search_result = vector_store.query(query_embedding, k=3)
print_result(query, search_result)

example.pyファイルを実行すると、出力は次のようになります。

Search result ("a swimming animal"):
- text: "fish", distance: 0.4562914811223072
- text: "dog", distance: 0.6469335836410557
- text: "tree", distance: 0.798545178640937

検索結果の 3 つの用語は、クエリされたベクトルからのそれぞれの距離によって並べ替えられます。距離が小さいほど、対応するdocument関連性が高くなります。

したがって、出力によれば、泳いでいる動物は魚、または泳ぐ才能のある犬である可能性が最も高いです。

参照

• ベクトルデータ型
• ベクター検索インデックス