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:
- Distância euclidiana: mede a distância mais curta entre dois vetores.
- Distância de cosseno: mede o cosseno do ângulo entre dois vetores.
- Vizinhos k-mais próximos (KNN): um algoritmo de aprendizado de máquina supervisionado usado para resolver problemas de classificação ou regressão.
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
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.Se você não tiver uma instância do Bigtable, siga as etapas em Criar uma instância.
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)
Realizar uma pesquisa KNN usando o GoogleSQL para Bigtable
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
- Crie aplicativos com tecnologia de LLM usando o LangChain.
- Saiba mais sobre como usar o SQL no Bigtable.