k 近傍法を探索して Bigtable で類似度ベクトル検索を行う

Preview — GoogleSQL for Bigtable

This feature is subject to the "Pre-GA Offerings Terms" in the General Service Terms section of the Service Specific Terms. Pre-GA features are available "as is" and might have limited support. For more information, see the launch stage descriptions.

類似性ベクトル検索を使用すると、Bigtable データ内の類似コンセプトとコンテキスト上の意味を特定できます。つまり、指定されたキー範囲内に保存されているデータをフィルタリングするときに、より関連性の高い結果を得ることができます。ユースケースの例としては、次のようなものがあります。

• 受信トレイ検索: 特定のユーザーのメールのセマンティック マッチングを実行する場合
• センサー範囲内の異常検出
• 検索拡張生成（RAG）用に既知のキーセット内で最も関連性の高いドキュメントを取得する

このページでは、GoogleSQL for Bigtable のコサイン距離ベクトル関数とユークリッド距離ベクトル関数を使用して k 近傍法を探索し、Bigtable で類似度ベクトル検索を行う方法について説明します。このページを読む前に、次のコンセプトを理解しておく必要があります。

• ユークリッド距離: 2 つのベクトル間の最短距離を測定します。
• コサイン距離: 2 つのベクトル間の角度のコサインを測定します。
• K 近傍法（KNN）: 分類や回帰の問題を解決するために使用される教師あり機械学習アルゴリズム。

Bigtable は、ベクトル エンベディングを操作する COSINE_DISTANCE() 関数と EUCLIDEAN_DISTANCE() 関数をサポートしています。これにより、入力エンベディングの KNN を探索できます。

Vertex AI Text Embeddings API を使用すると、Bigtable データをベクトル エンベディングとして生成して保存できます。これらのベクトル エンベディングをクエリの入力パラメータとして指定して、N 次元空間内で最も近いベクトルを探索し、意味的に類似または関連するアイテムを検索できます。

どちらの距離関数も、タイプ array<> の引数 vector1 と vector2 を使用し、ディメンションと長さが同じである必要があります。これらの関数の詳細については、以下をご覧ください。

• COSINE_DISTANCE()
• EUCLIDEAN_DISTANCE()

このページのコードは、エンベディングを作成して Bigtable に保存し、KNN 検索を実行する方法を示しています。

このページの例では、EUCLIDEAN_DISTANCE() と Python 用 Bigtable クライアント ライブラリを使用します。ただし、COSINE_DISTANCE() と、Bigtable 用 GoogleSQL をサポートする任意のクライアント ライブラリ（ Java 用 Bigtable クライアント ライブラリなど）を使用することもできます。

始める前に

コードサンプルを試す前に、次の操作を行います。

必要なロール

Bigtable への読み取りと書き込みに必要な権限を取得するには、次の IAM ロールを付与するよう管理者に依頼してください。

• リクエストを送信する Bigtable インスタンスの Bigtable ユーザー（roles/bigtable.user）

環境の設定

1. Python 用 Bigtable クライアント ライブラリをダウンロードしてインストールします。GoogleSQL for Bigtable 関数を使用するには、python-bigtable バージョン 2.26.0 以降を使用する必要があります。認証の設定方法など、手順については、Python hello world をご覧ください。

2. Bigtable インスタンスがない場合は、インスタンスを作成するの手順に沿って作成します。

3. リソース ID を特定します。コードを実行するときに、次のプレースホルダを Google Cloud プロジェクト、Bigtable インスタンス、テーブルの ID に置き換えます。

   • PROJECT_ID
   • INSTANCE_ID
   • TABLE_ID

テキスト、エンベディング、検索フレーズを保存するテーブルを作成する

2 つの列ファミリーを持つテーブルを作成します。

Python

from google.cloud import bigtable
from google.cloud.bigtable import column_family

client = bigtable.Client(project=PROJECT_ID, admin=True)
instance = client.instance(INSTANCE_ID)
table = instance.table(TABLE_ID)
column_families = {"docs":column_family.MaxVersionsGCRule(2), "search_phrase":column_family.MaxVersionsGCRule(2)}

if not table.exists():
  table.create(column_families=column_families)
else:
  print("Table already exists")

Vertex の事前トレーニング済みの基盤モデルを使用してテキストをエンベディングする

テキストとエンベディングを生成し、関連するキーとともに Bigtable に保存します。その他のドキュメントについては、テキスト エンベディングを取得するまたはマルチモーダル エンベディングを取得するをご覧ください。

Python

from typing import List, Optional
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
from vertexai.generative_models import GenerativeModel

#defines which LLM that we should use to generate the text
model = GenerativeModel("gemini-1.5-pro-001")

#First, use generative AI to create a list of 10 chunks for phrases
#This can be replaced with a static list of text items or your own data

chunks = []
for i in range(10):
  response = model.generate_content(
      "Generate a paragraph between 10 and 20 words that is about about either
      Bigtable or Generative AI"
)
chunks.append(response.text)
print(response.text)
#create embeddings for the chunks of text
def embed_text(
  texts: List[str] = chunks,
  task: str = "RETRIEVAL_DOCUMENT",
  model_name: str = "text-embedding-004",
  dimensionality: Optional[int] = 128,
) -> List[List[float]]:
  """Embeds texts with a pre-trained, foundational model."""
  model = TextEmbeddingModel.from_pretrained(model_name)
  inputs = [TextEmbeddingInput(text, task) for text in texts]
  kwargs = dict(output_dimensionality=dimensionality) if dimensionality else {}
  embeddings = model.get_embeddings(inputs, **kwargs)
  return [embedding.values for embedding in embeddings]

embeddings = embed_text()
print("embeddings created for text phrases")

バイト オブジェクトに変換できる関数を定義する

Bigtable は Key-Value ペア用に最適化されており、通常はデータをバイト オブジェクトとして保存します。Bigtable 用のデータモデルの設計の詳細については、スキーマ設計のベスト プラクティスをご覧ください。

Vertex から返されたエンベディングは、Python で浮動小数点数のリストとして保存されているため、変換する必要があります。各要素をビッグエンディアン IEEE 754 浮動小数点数に変換してから、連結します。次の関数でこの処理を行います。

Python

import struct
def floats_to_bytes(float_list):
  """
  Convert a list of floats to a bytes object, where each float is represented
  by 4 big-endian bytes.

  Parameters:
  float_list (list of float): The list of floats to be converted.

  Returns:
  bytes: The resulting bytes object with concatenated 4-byte big-endian
  representations of the floats.
  """
  byte_array = bytearray()

  for value in float_list:
      packed_value = struct.pack('>f', value)
      byte_array.extend(packed_value)

  # Convert bytearray to bytes
  return bytes(byte_array)

エンベディングを Bigtable に書き込む

エンベディングをバイト オブジェクトに変換し、ミューテーションを作成して、データを Bigtable に書き込みます。

Python

from google.cloud.bigtable.data  import RowMutationEntry
from google.cloud.bigtable.data  import SetCell

mutations = []
embeddings = embed_text()
for i, embedding in enumerate(embeddings):
  print(embedding)

  #convert each embedding into a byte object
  vector = floats_to_bytes(embedding)

  #set the row key which will be used to pull the range of documents (ex. doc type or user id)
  row_key = f"doc_{i}"

  row = table.direct_row(row_key)

  #set the column for the embedding based on the byte object format of the embedding
  row.set_cell("docs","embedding",vector)
  #store the text associated with vector in the same key
  row.set_cell("docs","text",chunks[i])
  mutations.append(row)

#write the rows to Bigtable
table.mutate_rows(mutations)

Bigtable 用の GoogleSQL を使用して KNN 検索を実行する

ベクトルはバイナリ エンコードされたデータとして保存され、BYTES 型から ARRAY<FLOAT32> への変換関数を使用して Bigtable から読み取ることができます。

SQL クエリは次のとおりです。

SELECT _key, TO_VECTOR32(data['embedding']) AS embedding
FROM table WHERE _key LIKE 'store123%';

Python では、GoogleSQL の COSINE_DISTANCE 関数を使用して、テキスト エンベディングと指定した検索フレーズとの類似度を確認できます。この計算には時間がかかるため、Python クライアント ライブラリの非同期データ クライアントを使用して SQL クエリを実行します。

Python

from google.cloud.bigtable.data import BigtableDataClientAsync

#first embed the search phrase
search_embedding = embed_text(texts=["Apache HBase"])

query = """
      select _key, docs['text'] as description
      FROM knn_intro
      ORDER BY COSINE_DISTANCE(TO_VECTOR32(docs['embedding']), {search_embedding})
      LIMIT 1;
      """

async def execute_query():
  async with BigtableDataClientAsync(project=PROJECT_ID) as client:
    local_query = query
    async for row in await client.execute_query(query.format(search_embedding=search_embedding[0]), INSTANCE_ID):
      return(row["_key"],row["description"])

await execute_query()

返されるレスポンスは、Bigtable を説明する生成されたテキストの説明です。

次のステップ

• LangChain を使用して、LLM を活用したアプリケーションを構築します。
• Bigtable での SQL の使用の詳細