Effectuer une recherche vectorielle de similarité dans Bigtable en recherchant les k plus proches voisins

La recherche vectorielle de similarité peut vous aider à identifier des concepts et des significations contextuelles similaires dans vos données Bigtable. Elle peut donc générer des résultats plus pertinents lorsque vous filtrez les données stockées dans une plage de clés spécifiée. Voici quelques exemples de cas d'utilisation:

  • Recherche dans la boîte de réception, où vous souhaitez effectuer une mise en correspondance sémantique des messages pour un utilisateur spécifique
  • Détection d'anomalies dans une plage de capteurs
  • Récupérer les documents les plus pertinents dans un ensemble de clés connues pour la génération augmentée par récupération (RAG)

Cette page explique comment effectuer une recherche vectorielle de similarité dans Bigtable à l'aide des fonctions vectorielles de distance cosinus et de distance euclidienne dans GoogleSQL pour Bigtable pour trouver les voisins les plus proches. Avant de lire cette page, vous devez connaître les concepts suivants:

Bigtable est compatible avec les fonctions COSINE_DISTANCE() et EUCLIDEAN_DISTANCE(), qui fonctionnent sur les représentations vectorielles continues, ce qui vous permet de trouver le KNN de l'embedding d'entrée.

Vous pouvez utiliser les API d'embeddings textuels Vertex AI pour générer et stocker vos données Bigtable sous forme d'embeddings vectoriels. Vous pouvez ensuite fournir ces représentations vectorielles continues en tant que paramètre d'entrée dans votre requête pour trouver les vecteurs les plus proches dans l'espace à N dimensions afin de rechercher des éléments sémantiquement similaires ou associés.

Les deux fonctions de distance acceptent les arguments vector1 et vector2, qui sont de type array<> et doivent avoir les mêmes dimensions et la même longueur. Pour en savoir plus sur ces fonctions, consultez les ressources suivantes:

Le code de cette page montre comment créer des représentations vectorielles continues, les stocker dans Bigtable, puis effectuer une recherche KNN.

L'exemple de cette page utilise EUCLIDEAN_DISTANCE() et la bibliothèque cliente Bigtable pour Python. Toutefois, vous pouvez également utiliser COSINE_DISTANCE() et n'importe quelle bibliothèque cliente compatible avec GoogleSQL pour Bigtable, comme la bibliothèque cliente Bigtable pour Java.

Avant de commencer

Effectuez les opérations suivantes avant d'essayer les exemples de code.

Rôles requis

Pour obtenir les autorisations nécessaires pour lire et écrire dans Bigtable, demandez à votre administrateur de vous accorder le rôle IAM suivant.

  • Utilisateur Bigtable (roles/bigtable.user) sur l'instance Bigtable à laquelle vous souhaitez envoyer des requêtes

Configurer votre environnement

  1. Téléchargez et installez la bibliothèque cliente Bigtable pour Python. Pour utiliser les fonctions GoogleSQL pour Bigtable, vous devez utiliser python-bigtable version 2.26.0 ou ultérieure. Pour obtenir des instructions, y compris sur la configuration de l'authentification, consultez Bonjour le monde en Python.

  2. Si vous ne disposez pas d'instance Bigtable, suivez la procédure décrite dans la section Créer une instance.

  3. Identifiez vos ID de ressource. Lorsque vous exécutez le code, remplacez les espaces réservés suivants par les ID de votre projet Google Cloud, de votre instance Bigtable et de votre table:

    • PROJECT_ID
    • INSTANCE_ID
    • TABLE_ID

Créer une table pour stocker le texte, les embeddings et la phrase de recherche

Créez une table avec deux familles de colonnes.

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

Générer des représentations vectorielles continues de texte avec un modèle de fondation pré-entraîné de Vertex

Générez le texte et les représentations vectorielles continues à stocker dans Bigtable, ainsi que les clés associées. Pour en savoir plus, consultez Obtenir des embeddings textuels ou Obtenir des embeddings multimodaux.

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

Définir des fonctions qui vous permettent de convertir en objets d'octets

Bigtable est optimisé pour les paires clé-valeur et stocke généralement les données en tant qu'objets d'octets. Pour en savoir plus sur la conception de votre modèle de données pour Bigtable, consultez la page Bonnes pratiques de conception de schémas.

Vous devez convertir les représentations vectorielles continues renvoyées par Vertex, qui sont stockées sous forme de liste de nombres à virgule flottante en Python. Vous convertissez chaque élément en format à virgule flottante IEEE 754 big-endian, puis vous les concatenates. La fonction suivante permet de le faire.

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)

Écrire les représentations vectorielles continues dans Bigtable

Convertissez les représentations vectorielles continues en objets octets, créez une mutation, puis écrivez les données dans 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)

Les vecteurs sont stockés sous forme de données encodées en binaire qui peuvent être lues à partir de Bigtable à l'aide d'une fonction de conversion du type BYTES au type ARRAY<FLOAT32>.

Voici la requête SQL:

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

En Python, vous pouvez utiliser la fonction COSINE_DISTANCE GoogleSQL pour trouver la similarité entre vos embeddings textuels et les expressions de recherche que vous lui fournissez. Étant donné que ce calcul peut prendre du temps à traiter, utilisez le client de données asynchrone de la bibliothèque cliente Python pour exécuter la requête 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 réponse renvoyée est une description textuelle générée qui décrit Bigtable.

Étape suivante