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

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

  • Correspondência semântica de mensagens para um usuário específico na pesquisa na caixa de entrada.
  • Detecção de anomalias em uma variedade de sensores.
  • Recuperar os documentos mais relevantes em um conjunto de chaves conhecidas para geração aumentada de recuperação (RAG).
  • Personalização dos resultados da pesquisa para melhorar a experiência do usuário, recuperando e classificando os resultados com base nos comandos e preferências históricas armazenadas pelo Bigtable.
  • Recuperação de conversas semelhantes para encontrar e mostrar conversas anteriores que são contextualmente semelhantes ao chat atual de um usuário para uma experiência mais personalizada.
  • Eliminação de duplicação de comandos para identificar comandos idênticos ou semanticamente semelhantes enviados pelo mesmo usuário e evitar o processamento redundante de IA.

Nesta página, descrevemos como realizar uma pesquisa de vetor de similaridade no Bigtable usando as funções de vetor de distância euclidiana e do cosseno no GoogleSQL para Bigtable (link em inglês) para encontrar os K vizinhos mais próximos. Antes de ler esta página, é importante entender os seguintes conceitos:

O Bigtable é compatível com as funções COSINE_DISTANCE() e EUCLIDEAN_DISTANCE(), que operam em embeddings de vetor, permitindo encontrar o KNN do embedding de entrada.

É possível usar as APIs de embeddings de texto da Vertex AI para gerar e armazenar seus dados do Bigtable como embeddings de vetor. Em seguida, forneça esses embeddings de vetor como um parâmetro de entrada na sua 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 usam os argumentos vector1 e vector2, que são do tipo array<> e precisam consistir nas mesmas dimensões e ter 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 de KNN.

O exemplo nesta página usa EUCLIDEAN_DISTANCE() e a biblioteca de cliente do Bigtable para Python. No entanto, também é possível 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 onde 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 o GoogleSQL com funções do Bigtable, use a versão 2.26.0 ou mais recente do python-bigtable. 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 seguintes marcadores de posição pelos IDs do seu projeto Google Cloud , da instância e da tabela do Bigtable:

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

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

Crie uma tabela com dois grupos de colunas.

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

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

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

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

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 de esquema.

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

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)

Gravar os embeddings no Bigtable

Converta os embeddings em objetos de byte, crie uma mutação e grave os dados no 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)

Os vetores são armazenados como dados codificados em formato binário que podem ser lidos do 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%';

Em Python, use a função COSINE_DISTANCE do GoogleSQL para encontrar a similaridade entre seus 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.

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

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

A seguir