Pesquisar com embeddings vetoriais

A página mostra como usar o Firestore para realizar pesquisas de vetores de vizinho mais próximo (KNN, na sigla em inglês) usando estas técnicas:

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

Armazenar embeddings de vetores

É possível criar valores vetoriais, como embeddings de texto, usando seus dados do Firestore e armazená-los em documentos do Firestore.

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

Veja no exemplo a seguir como armazenar um embedding de vetor em um documento do Firestore:

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

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 vetoriais 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 pesquisa de vizinho mais próximo com seus embeddings vetoriais, é preciso criar um índice correspondente. Os exemplos a seguir demonstram como criar e gerenciar índices de vetores.

Criar um índice de vetor de campo único

Para criar um índice de vetor de campo único, 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
    

onde:

  • collection-group é o ID do grupo de coleções.
  • vector-field é o nome do campo que contém o embedding vetorial.
  • 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": "{}"}.

Criar um índice vetorial composto

O exemplo a seguir cria um índice de vetor composto para o campo color e um campo de embedding de vetor.

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=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
    

onde:

  • 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
    

onde:

  • 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 um índice não existir, o Firestore vai sugerir um índice para ser criado usando a CLI gcloud.

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 vizinhas mais próximas são compatíveis com as seguintes opções de distância vetorial:

  • EUCLIDEAN: mede a distância EUCLIDEAN entre os vetores. Para saber mais, consulte Euclidiano (link em inglês).
  • COSINE: compara vetores com base no ângulo entre eles, o que permite medir a semelhança não baseada na magnitude dos vetores. Recomendamos o uso de DOT_PRODUCT com vetores normalizados por unidade em vez da distância COSINE, que é matematicamente equivalente a um desempenho melhor. Para saber mais, consulte Semelhança de cossenos (link em inglês).
  • DOT_PRODUCT: semelhante a COSINE, mas é afetado pela magnitude dos 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 os filtros de desigualdade. Os filtros compostos and e or são compatíveis. Os seguintes filtros são compatíveis com filtros de campo:

  • == 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 vetoriais, observe as seguintes limitações:

  • A dimensão de incorporação máxima aceita é 2048. Para armazenar índices maiores, use a redução de dimensionalidade.
  • O número máximo de documentos a serem retornados de uma consulta vizinha mais próxima é 1.000.
  • A pesquisa 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 suporte à pesquisa de vetor.

A seguir