通过查找 K 个最近邻，在 Bigtable 中执行相似向量搜索

相似度向量搜索可帮助您识别 Bigtable 数据中的相似概念和上下文含义，这意味着，在过滤存储在指定键范围内的数据时，它可以生成更相关的结果。示例用例可能包括：

  • 收件箱搜索，您希望针对特定用户执行邮件语义匹配
  • 在一系列传感器内检测异常
  • 从一组已知键中检索最相关的文档，以进行检索增强生成 (RAG)

本页介绍了如何使用 GoogleSQL for Bigtable 中的余弦距离和欧几里得距离向量函数在 Bigtable 中执行相似度向量搜索，以查找 K 个最近邻。在阅读本页内容之前，请务必了解以下概念：

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 角色。

设置环境

  1. 下载并安装 Python 版 Bigtable 客户端库。 如需使用适用于 Bigtable 的 GoogleSQL 函数，您必须使用 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 的生成文本说明。

后续步骤