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:
- Distância euclidiana: mede a menor distância entre dois vetores.
- Distância do 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 é 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
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.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 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)
Realizar uma pesquisa de KNN usando o GoogleSQL para Bigtable
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
- Criar aplicativos com tecnologia de LLM usando o LangChain.
- Saiba mais sobre como usar o SQL no Bigtable.