Pesquisar com embeddings de vetor

A página mostra como usar o Cloud Firestore para fazer pesquisas de vetor de vizinho mais próximo (KNN) usando estas técnicas:

  • Armazenar valores vetoriais
  • Criar e gerenciar índices de vetor KNN
  • Fazer uma consulta de vizinho mais próximo (KNN) usando uma das funções de distância de vetor compatíveis

Armazenar embeddings de vetores

Você pode criar valores vetoriais, como embeddings de texto, dos seus dados do Cloud Firestore e armazená-los em documentos do Cloud Firestore.

Operação de gravação com um embedding de vetor

O exemplo a seguir mostra como armazenar um embedding de vetor em um documento do Cloud Firestore:

Python
from google.cloud import firestore
from google.cloud.firestore_v1.vector import Vector

firestore_client = firestore.Client()
collection = firestore_client.collection("coffee-beans")
doc = {
  "name": "Kahawa coffee beans",
  "description": "Information about the Kahawa coffee beans.",
  "embedding_field": Vector([1.0 , 2.0, 3.0])
}

collection.add(doc)
    
Node.js
import {
  Firestore,
  FieldValue,
} from "@google-cloud/firestore";

const db = new Firestore();
const coll = db.collection('coffee-beans');
await coll.add({
  name: "Kahawa coffee beans",
  description: "Information about the Kahawa coffee beans.",
  embedding_field: FieldValue.vector([1.0 , 2.0, 3.0])
});
    

Calcular embeddings de vetores com uma função do Cloud

Para calcular e armazenar embeddings de vetor sempre que um documento for atualizado ou criado, configure uma função do Cloud:

Python
@functions_framework.cloud_event
def store_embedding(cloud_event) -> None:
  """Triggers by a change to a Firestore document.
  """
  firestore_payload = firestore.DocumentEventData()
  payload = firestore_payload._pb.ParseFromString(cloud_event.data)

  collection_id, doc_id = from_payload(payload)
  # Call a function to calculate the embedding
  embedding = calculate_embedding(payload)
  # Update the document
  doc = firestore_client.collection(collection_id).document(doc_id)
  doc.set({"embedding_field": embedding}, merge=True)
    
Node.js
/**
 * A vector embedding will be computed from the
 * value of the `content` field. The vector value
 * will be stored in the `embedding` field. The
 * field names `content` and `embedding` are arbitrary
 * field names chosen for this example.
 */
async function storeEmbedding(event: FirestoreEvent<any>): Promise<void> {
  // Get the previous value of the document's `content` field.
  const previousDocumentSnapshot = event.data.before as QueryDocumentSnapshot;
  const previousContent = previousDocumentSnapshot.get("content");

  // Get the current value of the document's `content` field.
  const currentDocumentSnapshot = event.data.after as QueryDocumentSnapshot;
  const currentContent = currentDocumentSnapshot.get("content");

  // Don't update the embedding if the content field did not change
  if (previousContent === currentContent) {
    return;
  }

  // Call a function to calculate the embedding for the value
  // of the `content` field.
  const embeddingVector = calculateEmbedding(currentContent);

  // Update the `embedding` field on the document.
  await currentDocumentSnapshot.ref.update({
    embedding: embeddingVector,
  });
}
    

Criar e gerenciar índices vetoriais

Antes de realizar a consulta de vizinho mais próximo com seus embeddings vetoriais, crie um índice correspondente. Os exemplos a seguir demonstram como criar e gerenciar índices vetoriais.

Criar um índice vetorial

Para criar um índice de vetores, use gcloud alpha firestore indexes composite create:

gcloud
gcloud alpha firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config field-path=vector-field,vector-config='vector-configuration' \
--database=database-id
    

em que:

  • collection-group é o ID do grupo de coleções.
  • vector-field é o nome do campo que contém o embedding de vetor.
  • database-id é o ID do banco de dados.
  • vector-configuration inclui o vetor dimension e o tipo de índice. O dimension é um número inteiro até 2.048. O tipo de índice precisa ser flat. Formate a configuração do índice da seguinte maneira: {"dimension":"DIMENSION", "flat": "{}"}.

O exemplo a seguir cria um índice composto, incluindo um índice de vetores para o campo vector-field e um índice crescente para o campo color. Você pode usar esse tipo de índice para pré-filtrar dados antes de uma pesquisa de vizinho mais próximo.

gcloud
gcloud alpha firestore indexes composite create \
--collection-group=collection-group \
--query-scope=COLLECTION \
--field-config=order=ASCENDING,field-path="color" \
--field-config field-path=vector-field,vector-config='{"dimension":"1024", "flat": "{}"}' \
--database=database-id
    

Listar todos os índices vetoriais

gcloud
gcloud alpha firestore indexes composite list --database=database-id

Substitua database-id pelo ID do banco de dados.

Excluir um índice vetorial

gcloud
gcloud alpha firestore indexes composite delete index-id --database=database-id
    

em que:

  • index-id é o ID do índice a ser excluído. Use indexes composite list para recuperar o ID do índice.
  • database-id é o ID do banco de dados.

Descrever um índice vetorial

gcloud
gcloud alpha firestore indexes composite describe index-id --database=database-id
    

em que:

  • index-id é o ID do índice a ser descrito. Use ou indexes composite list para recuperar o ID do índice.
  • database-id é o ID do banco de dados.

Fazer uma consulta de vizinho mais próximo

É possível realizar uma pesquisa por similaridade para encontrar os vizinhos mais próximos de um embedding de vetor. As pesquisas por similaridade exigem índices vetoriais. Se não houver um índice, o Cloud Firestore vai sugerir a criação usando a gcloud CLI.

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure

collection = collection("coffee-beans")

# Requires vector index
collection.find_nearest(
   vector_field="embedding_field",
   query_vector=Vector([3.0, 1.0, 2.0]),
   distance_measure=DistanceMeasure.EUCLIDEAN,
   limit=5)
    
Node.js
import {
  Firestore,
  FieldValue,
  VectorQuery,
  VectorQuerySnapshot,
} from "@google-cloud/firestore";

// Requires single-field vector index
const vectorQuery: VectorQuery = coll.findNearest('embedding_field', FieldValue.vector([3.0, 1.0, 2.0]), {
  limit: 5,
  distanceMeasure: 'EUCLIDEAN'
});

const vectorQuerySnapshot: VectorQuerySnapshot = await vectorQuery.get();
    

Distâncias vetoriais

As consultas de vizinhos mais próximos são compatíveis com as seguintes opções de distância vetorial:

  • EUCLIDEAN: mede a distância EUCLIDIANA entre os vetores. Para saber mais, consulte Euclidiano.
  • COSINE: compara vetores com base no ângulo entre eles, o que permite medir a similaridade que não se baseia na magnitude dos vetores. Recomendamos usar DOT_PRODUCT com vetores normalizados unitários em vez da distância COSSENO, que é matematicamente equivalente a um melhor desempenho. Para saber mais, consulte Semelhança de cossenos.
  • DOT_PRODUCT: semelhante a COSINE, mas é afetado pela magnitude de vetores. Para saber mais, consulte Produto escalar.

Pré-filtrar dados

Para pré-filtrar os dados antes de encontrar os vizinhos mais próximos, é possível combinar uma pesquisa por similaridade com outros filtros, exceto filtros de desigualdade. Os filtros compostos and e or são compatíveis. Para filtros de campo, os seguintes filtros são compatíveis:

  • == igual a
  • in
  • array_contains
  • array_contains_any
Python
# Similarity search with pre-filter
# Requires composite vector index
collection.where("color", "==", "red").find_nearest(
   vector_field="embedding_field",
   query_vector=Vector([3.0, 1.0, 2.0]),
   distance_measure=DistanceMeasure.EUCLIDEAN,
   limit=5)
    
Node.js
// Similarity search with pre-filter
// Requires composite vector index
const preFilteredVectorQuery: VectorQuery = coll
  .where("color", "==", "red")
  .findNearest("embedding_field", FieldValue.vector([3.0, 1.0, 2.0]), {
    limit: 5,
    distanceMeasure: "EUCLIDEAN",
  });

vectorQueryResults = await preFilteredVectorQuery.get();
    

Limitações

Ao trabalhar com embeddings de vetor, observe as seguintes limitações:

  • A dimensão de incorporação máxima compatível é 2048. Para armazenar índices maiores, use redução de dimensionalidade.
  • O número máximo de documentos a serem retornados de uma consulta de vizinho mais próximo é 1.000.
  • A consulta de vetor não é compatível com listeners de snapshots em tempo real.
  • Não é possível usar filtros de desigualdade para pré-filtrar dados.
  • Somente as bibliotecas de cliente Python e Node.js oferecem compatibilidade com a pesquisa de vetor.

A seguir