Realizar búsquedas de vectores de similitud en Bigtable buscando los K vecinos más cercanos

La búsqueda de vectores de similitud puede ayudarte a identificar conceptos similares y significados contextuales en tus datos de Bigtable, lo que significa que puede proporcionar resultados más relevantes al filtrar los datos almacenados en un intervalo de claves específico. Estos son algunos ejemplos de casos prácticos:

  • Concordancia semántica de los mensajes de un usuario concreto en la búsqueda de la bandeja de entrada.
  • Detección de anomalías en una serie de sensores.
  • Recupera los documentos más relevantes de un conjunto de claves conocidas para la generación aumentada por recuperación (RAG).
  • Personalización de los resultados de búsqueda para mejorar la experiencia de búsqueda de un usuario. Para ello, se obtienen y se clasifican los resultados en función de sus peticiones e historial y preferencias, que se almacenan en Bigtable.
  • Recuperación de conversaciones similares para encontrar y mostrar conversaciones anteriores que sean contextualmente similares a la conversación actual de un usuario para ofrecer una experiencia más personalizada.
  • Anula los duplicados de las peticiones para identificar las que sean idénticas o semánticamente similares enviadas por el mismo usuario y evitar que la IA las procese de forma redundante.

En esta página se describe cómo realizar búsquedas de vectores de similitud en Bigtable mediante las funciones de vectores de distancia euclídea y de distancia del coseno en GoogleSQL para Bigtable para encontrar los K vecinos más cercanos. Antes de leer esta página, es importante que conozcas los siguientes conceptos:

Bigtable admite las funciones COSINE_DISTANCE() y EUCLIDEAN_DISTANCE(), que operan en incrustaciones de vectores, lo que te permite encontrar los vecinos más cercanos (KNN) de la incrustación de entrada.

Puedes usar las APIs de inserciones de texto de Vertex AI para generar y almacenar tus datos de Bigtable como inserciones de vectores. Después, puedes proporcionar estas incrustaciones de vectores como parámetro de entrada en tu consulta para encontrar los vectores más cercanos en el espacio N-dimensional y buscar elementos semánticamente similares o relacionados.

Ambas funciones de distancia toman los argumentos vector1 y vector2, que son del tipo array<> y deben tener las mismas dimensiones y la misma longitud. Para obtener más información sobre estas funciones, consulta los siguientes artículos:

El código de esta página muestra cómo crear inserciones, almacenarlas en Bigtable y, a continuación, realizar una búsqueda de los k vecinos más cercanos.

En el ejemplo de esta página se usa EUCLIDEAN_DISTANCE() y la biblioteca de cliente de Bigtable para Python. Sin embargo, también puedes usar COSINE_DISTANCE() y cualquier biblioteca de cliente que admita GoogleSQL para Bigtable, como la biblioteca de cliente de Bigtable para Java.

Antes de empezar

Completa los siguientes pasos antes de probar los ejemplos de código.

Roles obligatorios

Para obtener los permisos que necesitas para leer y escribir en Bigtable, pide a tu administrador que te conceda el siguiente rol de gestión de identidades y accesos:

  • Usuario de Bigtable (roles/bigtable.user) en la instancia de Bigtable a la que quieras enviar solicitudes

Configurar un entorno

  1. Descarga e instala la biblioteca de cliente de Bigtable para Python. Para usar las funciones de GoogleSQL para Bigtable, debes usar la versión python-bigtable 2.26.0 o una posterior. Las instrucciones, incluido cómo configurar la autenticación, se encuentran en Python hello world.

  2. Si no tienes ninguna instancia de Bigtable, sigue los pasos que se indican en Crear una instancia.

  3. Identifica los IDs de tus recursos. Cuando ejecutes el código, sustituye los siguientes marcadores de posición por los IDs de tu Google Cloud proyecto, instancia de Bigtable y tabla:

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

Crea una tabla para almacenar el texto, las incrustaciones y la frase de búsqueda

Crea una tabla con dos familias de columnas.

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

Insertar textos con un modelo fundacional preentrenado de Vertex

Genera el texto y las inserciones que se almacenarán en Bigtable junto con las claves asociadas. Para obtener más documentación, consulta Obtener embeddings de texto o Obtener embeddings multimodales.

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 funciones que te permitan convertirte en objetos de bytes

Bigtable está optimizado para pares clave-valor y, por lo general, almacena datos como objetos de bytes. Para obtener más información sobre cómo diseñar tu modelo de datos para Bigtable, consulta las prácticas recomendadas para el diseño de esquemas.

Debes convertir las inserciones que devuelve Vertex, que se almacenan como una lista de números de coma flotante en Python. Convierte cada elemento al formato de coma flotante IEEE 754 big-endian y, a continuación, concaténalos. La siguiente función consigue esto.

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)

Escribir las inserciones en Bigtable

Convierte las inserciones en objetos de bytes, crea una mutación y, a continuación, escribe los datos en 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)

Los vectores se almacenan como datos codificados en formato binario que se pueden leer de Bigtable mediante una función de conversión del tipo BYTES a ARRAY<FLOAT32>.

Esta es la consulta de SQL:

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

En Python, puedes usar la función COSINE_DISTANCE de GoogleSQL para encontrar la similitud entre tus inserciones de texto y las frases de búsqueda que le proporciones. Como este cálculo puede tardar en procesarse, usa el cliente de datos asíncrono de la biblioteca de cliente de Python para ejecutar la consulta de 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()

La respuesta que se devuelve es una descripción de texto generada que describe Bigtable.

Siguientes pasos