TiDB ベクトル検索を Jina AI 埋め込み API と統合する

このチュートリアルでは、 ジナ・アイ使用してテキスト データの埋め込みを生成し、その埋め込みを TiDB ベクトルstorageに保存して、埋め込みに基づいて類似のテキストを検索する方法について説明します。

注記：

ベクトル検索機能は、TiDB セルフマネージド クラスターとTiDB Cloudサーバーレスクラスターでのみ使用できます。

前提条件

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

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

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

• ローカルテストTiDBクラスタをデプロイまたは本番のTiDBクラスタをデプロイに従ってローカル クラスターを作成します。
• TiDB Cloud Serverless クラスターの作成に従って、独自のTiDB Cloudクラスターを作成します。

サンプルアプリを実行する

以下の手順に従って、TiDB Vector Search を JinaAI Embedding と統合する方法を簡単に学習できます。

ステップ1. リポジトリをクローンする

tidb-vector-pythonリポジトリをローカル マシンにクローンします。

git clone https://github.com/pingcap/tidb-vector-python.git

ステップ2. 仮想環境を作成する

プロジェクト用の仮想環境を作成します。

cd tidb-vector-python/examples/jina-ai-embeddings-demo
python3 -m venv .venv
source .venv/bin/activate

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

デモ プロジェクトに必要な依存関係をインストールします。

pip install -r requirements.txt

ステップ4. 環境変数を設定する

Jina AI 埋め込み APIページから Jina AI API キーを取得し、選択した TiDB デプロイメント オプションに応じて環境変数を構成します。

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

ステップ5.デモを実行する

python jina-ai-embeddings-demo.py

出力例:

- Inserting Data to TiDB...
  - Inserting: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - Inserting: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- List All Documents and Their Distances to the Query:
  - distance: 0.3585317326132522
    content: Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.
- The Most Relevant Document and Its Distance to the Query:
  - distance: 0.10858102967720984
    content: TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.

サンプルコードスニペット

Jina AIから埋め込みを取得する

Jina AI 埋め込み API を呼び出すgenerate_embeddingsヘルパー関数を定義します。

import os
import requests
import dotenv

dotenv.load_dotenv()

JINAAI_API_KEY = os.getenv('JINAAI_API_KEY')

def generate_embeddings(text: str):
    JINAAI_API_URL = 'https://api.jina.ai/v1/embeddings'
    JINAAI_HEADERS = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {JINAAI_API_KEY}'
    }
    JINAAI_REQUEST_DATA = {
        'input': [text],
        'model': 'jina-embeddings-v2-base-en'  # with dimension 768.
    }
    response = requests.post(JINAAI_API_URL, headers=JINAAI_HEADERS, json=JINAAI_REQUEST_DATA)
    return response.json()['data'][0]['embedding']

TiDBクラスターに接続する

SQLAlchemy を介して TiDB クラスターに接続します。

import os
import dotenv

from tidb_vector.sqlalchemy import VectorType
from sqlalchemy.orm import Session, declarative_base

dotenv.load_dotenv()

TIDB_DATABASE_URL = os.getenv('TIDB_DATABASE_URL')
assert TIDB_DATABASE_URL is not None
engine = create_engine(url=TIDB_DATABASE_URL, pool_recycle=300)

ベクターテーブルスキーマを定義する

テキストを格納するためのcontent列と埋め込みを格納するためのcontent_vecという名前のベクトル列を持つjinaai_tidb_demo_documentsという名前のテーブルを作成します。

from sqlalchemy import Column, Integer, String, create_engine
from sqlalchemy.orm import declarative_base

Base = declarative_base()

class Document(Base):
    __tablename__ = "jinaai_tidb_demo_documents"

    id = Column(Integer, primary_key=True)
    content = Column(String(255), nullable=False)
    content_vec = Column(
        # DIMENSIONS is determined by the embedding model,
        # for Jina AI's jina-embeddings-v2-base-en model it's 768.
        VectorType(dim=768),
        comment="hnsw(distance=cosine)"

注記：

• ベクトル列の次元は、埋め込みモデルによって生成された埋め込みの次元と一致する必要があります。
• この例では、 jina-embeddings-v2-base-enモデルによって生成される埋め込みの次元は768です。

Jina AIで埋め込みを作成し、TiDBに保存する

Jina AI Embeddings API を使用して、各テキストの埋め込みを生成し、その埋め込みを TiDB に保存します。

TEXTS = [
   'Jina AI offers best-in-class embeddings, reranker and prompt optimizer, enabling advanced multimodal AI.',
   'TiDB is an open-source MySQL-compatible database that supports Hybrid Transactional and Analytical Processing (HTAP) workloads.',
]
data = []

for text in TEXTS:
    # Generate embeddings for the texts via Jina AI API.
    embedding = generate_embeddings(text)
    data.append({
        'text': text,
        'embedding': embedding
    })

with Session(engine) as session:
   print('- Inserting Data to TiDB...')
   for item in data:
      print(f'  - Inserting: {item["text"]}')
      session.add(Document(
         content=item['text'],
         content_vec=item['embedding']
      ))
   session.commit()

TiDB で Jina AI 埋め込みを使用してセマンティック検索を実行する

Jina AI 埋め込み API を使用してクエリ テキストの埋め込みを生成し、クエリ テキストの埋め込みとベクトル テーブル内の各埋め込み間のコサイン距離に基づいて最も関連性の高いドキュメントを検索します。

query = 'What is TiDB?'
# Generate the embedding for the query via Jina AI API.
query_embedding = generate_embeddings(query)

with Session(engine) as session:
    print('- The Most Relevant Document and Its Distance to the Query:')
    doc, distance = session.query(
        Document,
        Document.content_vec.cosine_distance(query_embedding).label('distance')
    ).order_by(
        'distance'
    ).limit(1).first()
    print(f'  - distance: {distance}\n'
          f'    content: {doc.content}')

参照

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