Realizar pesquisa de vetor de similaridade no Bigtable encontrando os vizinhos mais próximos de K

A pesquisa por similaridade vetorial pode ajudar a identificar conceitos semelhantes e significado contextual nos dados do Bigtable, o que significa que ela pode gerar resultados mais relevantes ao filtrar dados armazenados em um intervalo de chave especificado. Exemplos de casos de uso:

  • Pesquisa na caixa de entrada, em que você quer fazer a correspondência semântica de mensagens para um usuário específico
  • Detecção de anomalias em uma série de sensores
  • Recuperar os documentos mais relevantes em um conjunto de chaves conhecidas para geração aumentada de recuperação (RAG, na sigla em inglês)

Esta página descreve como realizar a pesquisa de vetor de similaridade no Bigtable usando as funções de vetor de distância do cosseno e da distância euclidiana no GoogleSQL para Bigtable para encontrar os vizinhos mais próximos de K. Antes de ler esta página, é importante que você entenda os seguintes conceitos:

O Bigtable oferece suporte às funções COSINE_DISTANCE() e EUCLIDEAN_DISTANCE(), que operam em embeddings de vetor, permitindo que você encontre a KNN do embedding de entrada.

É possível usar as APIs de embedding de texto da Vertex AI para gerar e armazenar os dados do Bigtable como embeddings de vetor. Em seguida, você pode fornecer esses embeddings de vetor como um parâmetro de entrada na consulta para encontrar os vetores mais próximos no espaço N-dimensional e pesquisar itens semanticamente semelhantes ou relacionados.

As duas funções de distância aceitam os argumentos vector1 e vector2, que são do tipo array<> e precisam ter as mesmas dimensões e o mesmo comprimento. Para mais detalhes sobre essas funções, consulte:

O código nesta página demonstra como criar embeddings, armazená-los no Bigtable e realizar uma pesquisa KNN.

O exemplo nesta página usa EUCLIDEAN_DISTANCE() e a biblioteca de cliente do Bigtable para Python. No entanto, você também pode usar COSINE_DISTANCE() e qualquer biblioteca de cliente que ofereça suporte ao GoogleSQL para Bigtable, como a biblioteca de cliente do Bigtable para Java.

Antes de começar

Conclua as etapas a seguir antes de testar os exemplos de código.

Funções exigidas

Para receber as permissões necessárias para ler e gravar no Bigtable, peça ao administrador para conceder a você o seguinte papel do IAM.

  • Usuário do Bigtable (roles/bigtable.user) na instância do Bigtable para a qual você quer enviar solicitações

Configurar o ambiente

  1. Faça o download e instale a biblioteca de cliente do Bigtable para Python. Para usar as funções do GoogleSQL para Bigtable, use python-bigtable na versão 2.26.0 ou mais recente. As instruções, incluindo como configurar a autenticação, estão em Python hello world.

  2. Se você não tiver uma instância do Bigtable, siga as etapas em Criar uma instância.

  3. Identifique os IDs dos recursos. Ao executar o código, substitua os marcadores de posição a seguir pelos IDs do seu projeto do Google Cloud, da instância do Bigtable e da tabela:

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

Crie uma tabela para armazenar o texto, os embeddings e a frase de pesquisa

Crie uma tabela com dois grupos de colunas.

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")

Incorporar textos com um modelo de base pré-treinado da Vertex

Gere o texto e os embeddings para armazenar no Bigtable com as chaves associadas. Para mais documentação, consulte Usar embeddings de texto ou Usar embeddings multimodais.

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")

Definir funções que permitem converter em objetos de byte

O Bigtable é otimizado para pares de chave-valor e geralmente armazena dados como objetos de byte. Para mais informações sobre como projetar seu modelo de dados para o Bigtable, consulte Práticas recomendadas de design do esquema.

Você precisa converter os embeddings que retornam do Vertex, que são armazenados como uma lista de números de ponto flutuante em Python. Você converte cada elemento em formato de ponto flutuante IEEE 754 big-endian e os concatena. A função a seguir faz isso.

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)

Gravar as embeddings no Bigtable

Converta os embeddings em objetos de byte, crie uma mutação e grave os dados no Bigtable.

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)

Os vetores são armazenados como dados codificados em binário que podem ser lidos no Bigtable usando uma função de conversão do tipo BYTES para ARRAY<FLOAT32>.

Esta é a consulta SQL:

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

No Python, é possível usar a função COSINE_DISTANCE do GoogleSQL para encontrar a semelhança entre os embeddings de texto e as frases de pesquisa fornecidas. Como essa computação pode levar tempo para ser processada, use o cliente de dados assíncrono da biblioteca de cliente do Python para executar a consulta SQL.

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()

A resposta retornada é uma descrição de texto gerada que descreve o Bigtable.

A seguir