Effectuer une recherche de similarité vectorielle dans Bigtable en trouvant les K voisins les plus proches

La recherche de similarités vectorielles peut vous aider à identifier des concepts et des significations contextuelles similaires dans vos données Bigtable. Cela signifie qu'elle peut fournir 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 :

  • Correspondance sémantique des messages pour un utilisateur spécifique dans la recherche dans la boîte de réception.
  • 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).
  • Personnalisation des résultats de recherche pour améliorer l'expérience de recherche d'un utilisateur en récupérant et en classant les résultats en fonction de ses requêtes et préférences historiques stockées dans Bigtable.
  • Récupération de fils de discussion similaires pour trouver et afficher les conversations passées qui sont contextuellement similaires à la conversation actuelle d'un utilisateur, afin de lui offrir une expérience plus personnalisée.
  • La déduplication des requêtes permet d'identifier les requêtes identiques ou sémantiquement similaires envoyées par le même utilisateur et d'éviter un traitement redondant par l'IA.

Cette page explique comment effectuer une recherche de similarité vectorielle dans Bigtable à l'aide des fonctions vectorielles de distance cosinus et de distance euclidienne dans GoogleSQL pour Bigtable afin de trouver les K voisins les plus proches. Avant de lire cette page, il est important que vous compreniez les concepts suivants :

Bigtable accepte les fonctions COSINE_DISTANCE() et EUCLIDEAN_DISTANCE(), qui fonctionnent sur les embeddings vectoriels et vous permettent de trouver le KNN de l'embedding d'entrée.

Vous pouvez utiliser les API Vertex AI d'embeddings textuels pour générer et stocker vos données Bigtable sous forme d'embeddings vectoriels. Vous pouvez ensuite fournir ces embeddings vectoriels comme paramètre d'entrée dans votre requête pour trouver les vecteurs les plus proches dans l'espace N-dimensionnel et 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 comporter les mêmes dimensions et avoir la même longueur. Pour en savoir plus sur ces fonctions, consultez les ressources suivantes :

Le code sur cette page montre comment créer des embeddings, 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 étapes 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 la version 2.26.0 ou ultérieure de python-bigtable. Les instructions, y compris la configuration de l'authentification, sont disponibles sur la page Hello World en Python.

  2. Si vous ne disposez pas d'instance Bigtable, suivez les étapes décrites dans Créer une instance.

  3. Identifiez vos ID de ressources. 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 l'expression 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")

Intégrer des textes avec un modèle de fondation pré-entraîné de Vertex

Générez le texte et les embeddings à stocker dans Bigtable, ainsi que les clés associées. Pour en savoir plus, consultez Obtenir des embeddings de texte 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 permettant de convertir en objets byte

Bigtable est optimisé pour les paires clé/valeur et stocke généralement les données sous forme d'objets byte. Pour en savoir plus sur la conception de votre modèle de données pour Bigtable, consultez Bonnes pratiques liées à la conception de schémas.

Vous devez convertir les embeddings renvoyés par Vertex, qui sont stockés sous forme de liste de nombres à virgule flottante en Python. Vous convertissez chaque élément au format à virgule flottante IEEE 754 big-endian, puis vous les concaténez. La fonction suivante permet d'y parvenir.

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 embeddings dans Bigtable

Convertissez les embeddings en objets byte, 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 au format binaire, qui peuvent être lues à partir de Bigtable à l'aide d'une fonction de conversion du type BYTES en 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 de GoogleSQL pour trouver la similarité entre vos embeddings de texte et les expressions de recherche que vous lui fournissez. Étant donné que ce calcul peut prendre du temps, 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.

Étapes suivantes