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

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

  • 在收件箱搜索中针对特定用户进行消息的语义匹配。
  • 在一定范围的传感器内检测异常情况。
  • 在已知密钥集中检索最相关的文档，以用于检索增强生成 (RAG)。
  • 搜索结果个性化，通过根据 Bigtable 存储的用户历史提示和偏好检索和排名结果，提升用户搜索体验。
  • 检索类似对话串，以查找并显示在情境上与用户当前对话类似的过往对话，从而提供更个性化的体验。
  • 提示去重，以识别同一用户提交的相同或语义相似的提示，避免冗余的 AI 处理。

本页介绍了如何使用 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 客户端库。 如需使用 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

创建一个表来存储文本、嵌入和搜索短语

创建一个包含两个列族的表。

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 的生成文本说明。

后续步骤