Realize uma pesquisa de vetores de similaridade no Bigtable encontrando os K vizinhos mais próximos

A pesquisa de vetores de semelhança pode ajudar a identificar conceitos semelhantes e o significado contextual nos seus dados do Bigtable, o que significa que pode fornecer resultados mais relevantes quando filtra dados armazenados num intervalo de chaves especificado. Seguem-se alguns exemplos de utilização:

  • Correspondência semântica de mensagens para um utilizador específico na pesquisa da caixa de entrada.
  • Deteção de anomalias num intervalo de sensores.
  • Obter os documentos mais relevantes num conjunto de chaves conhecidas para a geração aumentada de obtenção (RAG).
  • Personalização dos resultados da pesquisa para melhorar a experiência de pesquisa de um utilizador através da obtenção e classificação de resultados com base nos respetivos comandos históricos e preferências armazenados pelo Bigtable.
  • Obtenção de discussões semelhantes para encontrar e apresentar conversas anteriores que sejam contextualmente semelhantes ao chat atual de um utilizador para uma experiência mais personalizada.
  • Solicitar a remoção de duplicados para identificar comandos idênticos ou semanticamente semelhantes enviados pelo mesmo utilizador e evitar o processamento de IA redundante.

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

O Bigtable suporta as funções COSINE_DISTANCE() e EUCLIDEAN_DISTANCE(), que operam em incorporações de vetores, permitindo-lhe encontrar o KNN da incorporação de entrada.

Pode usar as APIs de incorporações de texto da Vertex AI para gerar e armazenar os seus dados do Bigtable como incorporações vetoriais. Em seguida, pode fornecer estas incorporações de vetores como um parâmetro de entrada na sua consulta para encontrar os vetores mais próximos no espaço N-dimensional para pesquisar itens semanticamente semelhantes ou relacionados.

Ambas as funções de distância usam os argumentos vector1 e vector2, que são do tipo array<> e têm de consistir nas mesmas dimensões e ter o mesmo comprimento. Para mais detalhes sobre estas funções, consulte o seguinte:

O código nesta página demonstra como criar incorporações, armazená-las no Bigtable e, em seguida, realizar uma pesquisa KNN.

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

Antes de começar

Conclua os seguintes passos antes de experimentar os exemplos de código.

Funções necessárias

Para receber as autorizações de que precisa para ler e escrever no Bigtable, peça ao seu administrador para lhe conceder a seguinte função de IAM:

  • Utilizador do Bigtable (roles/bigtable.user) na instância do Bigtable para a qual quer enviar pedidos

Configure o seu ambiente

  1. Transfira e instale a biblioteca cliente do Bigtable para Python. Para usar as funções do GoogleSQL para Bigtable, tem de usar a versão python-bigtable 2.26.0 ou posterior. Pode encontrar instruções, incluindo como configurar a autenticação, em Python hello world.

  2. Se não tiver uma instância do Bigtable, siga os passos em Criar uma instância.

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

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

Crie uma tabela para armazenar o texto, as incorporações e a expressão de pesquisa

Crie uma tabela com duas famílias 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")

Incorpore textos com um modelo fundamental pré-preparado do Vertex

Gere o texto e as incorporações para armazenar no Bigtable, juntamente com as chaves associadas. Para ver documentação adicional, consulte os artigos Obtenha incorporações de texto ou Obtenha incorporações 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")

Defina funções que lhe permitem converter em objetos de bytes

O Bigtable está otimizado para pares de chave-valor e, geralmente, armazena dados como objetos de bytes. Para mais informações sobre a criação do modelo de dados para o Bigtable, consulte as práticas recomendadas de criação de esquemas.

Tem de converter as incorporações devolvidas pelo Vertex, que são armazenadas como uma lista de números de vírgula flutuante em Python. Converte cada elemento na formação de ponto flutuante IEEE 754 big-endian e, em seguida, concatena-os. A seguinte função consegue isto.

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)

Escreva as incorporações no Bigtable

Converta as incorporações em objetos de bytes, crie uma mutação e, em seguida, escreva 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 a partir do Bigtable através de uma função de conversão do tipo BYTES para ARRAY<FLOAT32>.

Aqui está a consulta SQL:

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

Em Python, pode usar a função COSINE_DISTANCE do GoogleSQL para encontrar a semelhança entre as suas incorporações de texto e as expressões de pesquisa que lhe fornece. Uma vez que este cálculo pode demorar algum tempo a ser processado, use o cliente de dados assíncronos da biblioteca cliente 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 devolvida é uma descrição de texto gerada que descreve o Bigtable.

O que se segue?