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

A pesquisa de vetor de 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. Os casos de uso de exemplo incluem o seguinte:

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

Antes de ler esta página, é importante entender 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 o KNN do embedding de entrada.

É possível usar as APIs de embeddings de texto da plataforma de agentes do Gemini Enterprise para gerar e armazenar seus 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 para 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 seguinte:

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 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, é necessário usar a versão 2.26.0 ou mais recente do python-bigtable. As instruções, incluindo como configurar a autenticação, estão em Hello world do Python.

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

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

    • 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 fundamental pré-treinado da plataforma de agentes

Gere o texto e os embeddings para armazenar no Bigtable junto com as chaves associadas. Para mais informações, 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 Bigtable, consulte Práticas recomendadas de design de esquema.

É necessário converter os embeddings que retornam da plataforma de agentes, que são armazenados como uma lista de números de ponto flutuante em Python. Converta cada elemento para a formação de usar pontos flutuantes IEEE 754 big-endian e, em seguida, concatene-os. A função a seguir realiza essa ação.

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 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%';

É possível usar a função COSINE_DISTANCE do GoogleSQL para encontrar a similaridade entre os embeddings de texto e as frases de pesquisa fornecidas. Como esse cálculo pode levar tempo para ser processado, 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