在 Bigtable 中執行相似度向量搜尋，找出 K 個最鄰近的項目

相似度向量搜尋可協助您在 Bigtable 資料中找出相似概念和情境意義，因此在篩選特定鍵範圍內儲存的資料時，可提供更相關的結果。應用實例包括：

  • 在收件匣搜尋中，對特定使用者的郵件進行語意比對。
  • 偵測感應器範圍內的異常狀況。
  • 在已知鍵的集合中，擷取最相關的文件，以進行檢索增強生成 (RAG)。
  • 個人化搜尋結果：根據 Bigtable 儲存的歷史提示和偏好設定，擷取並排序結果，提升使用者的搜尋體驗。
  • 系統會擷取類似的對話串，找出並顯示與使用者目前對話情境相似的過往對話，提供更個人化的體驗。
  • 系統會簡化提示，找出同一位使用者提交的相同或語意相似提示，避免 AI 重複處理。

本頁說明如何使用 Bigtable 適用的 GoogleSQL 中的餘弦距離和歐氏距離向量函式，在 Bigtable 中執行相似度向量搜尋，找出 K 個最鄰近的鄰居。閱讀本頁面之前，請務必瞭解下列概念：

  • 歐氏距離：測量兩個向量之間的最短距離。
  • 餘弦距離： 測量兩個向量間的角度餘弦值。
  • K 最鄰近 (KNN)： 一種監督式機器學習演算法，用於解決分類或迴歸問題。

Bigtable 支援 COSINE_DISTANCE() 和 EUCLIDEAN_DISTANCE() 函式，這些函式可對向量嵌入內容執行作業，讓您找出輸入嵌入內容的 KNN。

您可以使用 Vertex AI 文字嵌入 API，以向量嵌入的形式生成及儲存 Bigtable 資料。然後，您可以在查詢中提供這些向量嵌入做為輸入參數，在 N 維空間中找出最鄰近的向量，搜尋語意相似或相關的項目。

這兩個距離函式都會採用 vector1 和 vector2 引數，這些引數屬於 array<> 型別，且必須包含相同維度並具有相同長度。如要進一步瞭解這些函式，請參閱下列說明：

本頁的程式碼示範如何建立嵌入、將其儲存在 Bigtable 中，然後執行 KNN 搜尋。

本頁面的範例使用 EUCLIDEAN_DISTANCE() 和 Python 適用的 Bigtable 用戶端程式庫。不過，您也可以使用 COSINE_DISTANCE()，以及任何支援 Bigtable GoogleSQL 的用戶端程式庫，例如 Java 適用的 Bigtable 用戶端程式庫。

事前準備

請先完成下列步驟，再試用程式碼範例。

必要的角色

如要取得讀取及寫入 Bigtable 資料所需的權限，請要求管理員授予下列 IAM 角色：

  • Bigtable 執行個體上的 Bigtable 使用者 (roles/bigtable.user)，您要將要求傳送至該執行個體

設定環境

  1. 下載並安裝 Python 專用的 Bigtable 用戶端程式庫。 如要使用 GoogleSQL 執行 Bigtable 函式，您必須使用 python-bigtable 2.26.0 以上版本。如需操作說明 (包括如何設定驗證)，請參閱 Python Hello World。

  2. 如果沒有 Bigtable 執行個體，請按照「建立執行個體」一文中的步驟操作。

  3. 找出資源 ID。執行程式碼時，請將下列預留位置替換為 Google Cloud 專案、Bigtable 執行個體和資料表的 ID：

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

建立資料表，用於儲存文字、嵌入和搜尋詞組

建立含有兩個資料欄系列的資料表。

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 專為鍵/值組合而設計，一般會將資料儲存為位元組物件。如要進一步瞭解如何設計 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)

向量會儲存為二進位編碼資料，可使用從 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。

後續步驟