Pesquisar com embeddings de vetor

A página mostra como usar o 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 medidas de distância de vetor compatíveis

Armazenar embeddings de vetores

Você pode criar valores vetoriais, como embeddings de texto, dos seus dados do Firestore e armazená-los em documentos do 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 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)
vector_search.py
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 Run:

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

Antes de criar um índice de vetor, faça upgrade para a versão mais recente da Google Cloud CLI:

gcloud components update

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

gcloud
gcloud 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 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 firestore indexes composite list --database=database-id

Substitua database-id pelo ID do banco de dados.

Excluir um índice vetorial

gcloud
gcloud 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 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 Firestore vai sugerir a criação usando a CLI gcloud.

O exemplo a seguir encontra 10 vizinhos mais próximos do vetor de consulta.

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

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

// Requires a single-field vector index
const vectorQuery: VectorQuery = coll.findNearest({
  vectorField: 'embedding_field',
  queryVector: [3.0, 1.0, 2.0],
  limit: 10,
  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.

Escolha a medida de distância

Dependendo de todos os embeddings de vetores serem normalizados, é possível determinar qual medida de distância usar para encontrar a medida de distância. Um embedding de vetor normalizado tem uma magnitude (comprimento) de exatamente 1,0.

Além disso, se você souber com qual medida de distância seu modelo foi treinado, use essa medida para calcular a distância entre o vetor e embeddings.

Dados normalizados

Se você tiver um conjunto de dados em que todas as embeddings de vetores forem normalizadas, as três medidas de distância vão fornecer os mesmos resultados de pesquisa semântica. Basicamente, embora cada a medida de distância retorna um valor diferente, esses valores são classificados da mesma maneira. Quando as inclusões são normalizadas, o DOT_PRODUCT geralmente é o mais eficiente em termos computacionais, mas a diferença é insignificante na maioria dos casos. No entanto, se seus aplicativo é altamente sensível ao desempenho, o DOT_PRODUCT pode ajudar com ajuste de desempenho.

Dados não normalizados

Se você tiver um conjunto de dados em que as embeddings de vetor não forem normalizadas, não será matematicamente correto usar DOT_PRODUCT como uma medida de distância, porque o produto escalar não mede distância. Dependendo sobre como os embeddings foram gerados e que tipo de pesquisa é preferível, a medida de distância COSINE ou EUCLIDEAN produz resultados de pesquisa subjetivamente melhores do que as outras medidas de distância. O experimento com COSINE ou EUCLIDEAN pode necessário para determinar qual é o melhor para seu caso de uso.

Não ter certeza se os dados são normalizados ou não

Se você não tiver certeza se os dados estão normalizados e quiser usar DOT_PRODUCT, recomendamos usar COSINE. COSINE é como DOT_PRODUCT com a normalização integrada. A distância medida usando COSINE varia de 0 a 2. Um resultado próximo de 0 indica que os vetores são muito semelhantes.

Pré-filtrar documentos

Para pré-filtrar documentos antes de encontrar os vizinhos mais próximos, é possível combinar uma por similaridade com outros operadores de consulta. Os filtros compostos and e or são compatíveis. Para mais informações sobre filtros de campo compatíveis, consulte Operadores de consulta.

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

# Similarity search with pre-filter
# Requires a composite vector index
vector_query = 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,
)
vector_search.py
Node.js
// Similarity search with pre-filter
// Requires composite vector index
const preFilteredVectorQuery: VectorQuery = coll
    .where("color", "==", "red")
    .findNearest({
      vectorField: "embedding_field",
      queryVector: [3.0, 1.0, 2.0],
      limit: 5,
      distanceMeasure: "EUCLIDEAN",
    });

const vectorQueryResults = await preFilteredVectorQuery.get();

Extrair a distância vetorial calculada

É possível extrair a distância vetorial calculada atribuindo um nome de propriedade de saída distance_result_field na consulta FindNearest, conforme mostrado no exemplo a seguir:

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

vector_query = collection.find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([3.0, 1.0, 2.0]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_result_field="vector_distance",
)

docs = vector_query.stream()

for doc in docs:
    print(f"{doc.id}, Distance: {doc.get('vector_distance')}")
vector_search.py
Node.js
const vectorQuery: VectorQuery = coll.findNearest(
    {
      vectorField: 'embedding_field',
      queryVector: [3.0, 1.0, 2.0],
      limit: 10,
      distanceMeasure: 'EUCLIDEAN',
      distanceResultField: 'vector_distance'
    });

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

snapshot.forEach((doc) => {
  console.log(doc.id, ' Distance: ', doc.get('vector_distance'));
});

Se você quiser usar uma máscara de campo para retornar um subconjunto de campos do documento com uma distanceResultField, inclua também o valor de distanceResultField na máscara de campo, conforme mostrado no exemplo a seguir:

Python
vector_query = collection.select(["color", "vector_distance"]).find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([3.0, 1.0, 2.0]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_result_field="vector_distance",
)
vector_search.py
Node.js
const vectorQuery: VectorQuery = coll
    .select('color', 'vector_distance')
    .findNearest({
      vectorField: 'embedding_field',
      queryVector: [3.0, 1.0, 2.0],
      limit: 10,
      distanceMeasure: 'EUCLIDEAN',
      distanceResultField: 'vector_distance'
    });

Especificar um limite de distância

É possível especificar um limite de similaridade que retorna apenas documentos dentro do limite. O comportamento do campo "Limite" depende da medida de distância à sua escolha:

  • As distâncias de EUCLIDEAN e COSINE limitam o limite aos documentos em que é menor ou igual ao limite especificado. Essa distância medidas diminuem à medida que os vetores se tornam mais semelhantes.
  • A distância de DOT_PRODUCT limita o limite aos documentos em que a distância é maior ou igual ao limite especificado. As distâncias do produto escalar aumentam à medida que os vetores se tornam mais semelhantes.

O exemplo a seguir mostra como especificar um limite de distância para retornar até 10 documentos mais próximos que estão, no máximo, a 4,5 unidades de distância usando a métrica de distância EUCLIDEAN:

Python
from google.cloud.firestore_v1.base_vector_query import DistanceMeasure
from google.cloud.firestore_v1.vector import Vector

collection = db.collection("coffee-beans")

vector_query = collection.find_nearest(
    vector_field="embedding_field",
    query_vector=Vector([3.0, 1.0, 2.0]),
    distance_measure=DistanceMeasure.EUCLIDEAN,
    limit=10,
    distance_threshold=4.5,
)

docs = vector_query.stream()

for doc in docs:
    print(f"{doc.id}")
vector_search.py
Node.js
const vectorQuery: VectorQuery = coll.findNearest({
  vectorField: 'embedding_field',
  queryVector: [3.0, 1.0, 2.0],
  limit: 10,
  distanceMeasure: 'EUCLIDEAN',
  distanceThreshold: 4.5
});

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

snapshot.forEach((doc) => {
  console.log(doc.id);
});

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.
  • Somente as bibliotecas de cliente Python e Node.js oferecem compatibilidade com a pesquisa de vetor.

A seguir