Similarity Search con Spanner y Vertex AI

1. Introducción

Los avances recientes en el aprendizaje profundo permitieron representar texto y otros datos de una manera que captura el significado semántico. Esto dio lugar a un nuevo enfoque de búsqueda, llamado búsqueda vectorial, que utiliza representaciones vectoriales de texto (conocidas como embeddings) para encontrar los documentos más relevantes para la búsqueda de un usuario. La búsqueda de vectores es preferible a la búsqueda tradicional para aplicaciones como la búsqueda de indumentaria, en la que los usuarios suelen buscar artículos por su descripción, estilo o contexto, en lugar de por los nombres exactos de los productos o las marcas. Podemos integrar la base de datos de Cloud Spanner con la Búsqueda vectorial para realizar la correlación de similitud de vectores. Con el uso conjunto de Spanner y la Búsqueda de vectores, los clientes pueden crear una integración potente que combine la disponibilidad, la confiabilidad y la escala de Spanner con las capacidades avanzadas de búsqueda de similitud de la Búsqueda de vectores de Vertex AI. Esta búsqueda se realiza comparando las incorporaciones de los elementos en el índice de Vector Search y devolviendo las coincidencias más similares.

Caso de uso

Imagina que eres un científico de datos en una tienda de moda que intenta mantenerse al día con las tendencias, las búsquedas de productos y las recomendaciones que cambian rápidamente. El desafío es que tienes recursos limitados y silos de datos. En esta entrada de blog, se muestra cómo implementar un caso de uso de recomendación de prendas de vestir con un enfoque de búsqueda por similitud en los datos de prendas de vestir.Se abordan los siguientes pasos:

Datos obtenidos de Spanner
Vectores generados para los datos de indumentaria con ML.PREDICT y almacenados en Spanner
Datos vectoriales de Spanner integrados en la Búsqueda de vectores con trabajos de flujo de datos y flujo de trabajo
Se realiza una búsqueda vectorial para encontrar coincidencias de similitud para la entrada del usuario.

Compilaremos una aplicación web de demostración para realizar búsquedas de indumentaria según el texto ingresado por el usuario. La aplicación permite a los usuarios buscar prendas ingresando una descripción de texto.

Índice de Spanner a Vector Search:

Los datos de la búsqueda de indumentaria se almacenan en Spanner. Invocaremos la API de Embeddings de Vertex AI en la construcción de ML.PREDICT directamente desde los datos de Spanner. Luego, aprovecharemos los trabajos de Dataflow y Workflow que suben de forma masiva estos datos (inventario y embeddings) a la Búsqueda de vectores de Vertex AI y actualizan el índice.

Ejecuta consultas del usuario en el índice:

Cuando un usuario ingresa una descripción de prendas, la app genera los embeddings en tiempo real con la API de Text Embeddings. Luego, se envía como entrada a la API de Vector Search para encontrar 10 descripciones de productos relevantes en el índice y mostrar la imagen correspondiente.

Descripción general de la arquitectura

La arquitectura de la aplicación de Spanner y Vector Search se muestra en el siguiente diagrama de 2 partes:

Índice de Spanner a Vector Search:

App cliente para ejecutar consultas del usuario en el índice:

Qué compilarás

De Spanner al índice vectorial:

Base de datos de Spanner para almacenar y administrar los datos de origen y los embeddings correspondientes
Es un trabajo de Workflow que sube datos de forma masiva (ID y embeddings) a la base de datos de la Búsqueda de vectores de Vertex AI.
Una API de Vector Search que se usa para encontrar descripciones de productos relevantes en el índice.

Ejecuta consultas del usuario en el índice:

Una aplicación web que permite a los usuarios ingresar descripciones de texto de prendas de vestir, realiza una búsqueda de similitud con el extremo del índice implementado y devuelve las prendas más cercanas a la entrada.

Cómo funciona

Cuando un usuario ingresa una descripción de texto de prendas de vestir, la aplicación web envía la descripción a la API de Vector Search. Luego, la API de Vector Search usa los embeddings de las descripciones de prendas para encontrar las descripciones de productos más relevantes del índice. Luego, se muestran al usuario las descripciones de los productos y las imágenes correspondientes. El flujo de trabajo general es el siguiente:

Genera embeddings para los datos almacenados en Spanner.
Exportar y subir embeddings a un índice de Vector Search
Consulta el índice de la Búsqueda de vectores para encontrar elementos similares realizando una búsqueda de vecinos más cercanos.

2. Requisitos

Un navegador, como Chrome o Firefox.
Un proyecto de Google Cloud con la facturación habilitada.

Antes de comenzar

En la página del selector de proyectos de la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud.
Asegúrate de que la facturación esté habilitada para tu proyecto de Cloud. Obtén información para verificar si la facturación está habilitada en un proyecto.
Asegúrate de que todas las APIs necesarias (Cloud Spanner, Vertex AI y Google Cloud Storage) estén habilitadas.
Usarás Cloud Shell, un entorno de línea de comandos que se ejecuta en Google Cloud y que viene precargado con gcloud. Consulta la documentación para ver los comandos y el uso de gcloud. Si tu proyecto no está configurado, usa el siguiente comando para hacerlo:

gcloud config set project <YOUR_PROJECT_ID>

Navega a la página de Cloud Spanner con tu proyecto de Google Cloud activo para comenzar.

3. Backend: Crea tu fuente de datos y tus incorporaciones de Spanner

En este caso de uso, la base de datos de Spanner contiene el inventario de prendas de vestir con las imágenes y descripciones correspondientes. Asegúrate de generar embeddings para la descripción del texto y almacenarlos en tu base de datos de Spanner como ARRAY<float64>.

Crea los datos de Spanner

Crea una instancia llamada "spanner-vertex" y una base de datos llamada "spanner-vertex-embeddings". Crea una tabla con el DDL:

CREATE TABLE
  apparels ( id NUMERIC,
    category STRING(100),
    sub_category STRING(50),
    uri STRING(200),
    content STRING(2000),
    embedding ARRAY<FLOAT64>
    )
PRIMARY KEY
  (id);

Inserta datos en la tabla con la instrucción INSERT de SQL

Las secuencias de comandos de inserción para los datos de muestra están disponibles aquí.

Crea un modelo de Text Embeddings

Esto es necesario para que podamos generar embeddings para el contenido de la entrada. A continuación, se muestra el DDL para el mismo:

CREATE MODEL text_embeddings INPUT(content STRING(MAX))
OUTPUT(
  embeddings
    STRUCT<
      statistics STRUCT<truncated BOOL, token_count FLOAT64>,
      values ARRAY<FLOAT64>>
)
REMOTE OPTIONS (
  endpoint = '//aiplatform.googleapis.com/projects/abis-345004/locations/us-central1/publishers/google/models/textembedding-gecko');

Genera embeddings de texto para los datos de origen

Crea una tabla para almacenar los embeddings y, luego, inserta los embeddings generados. En una aplicación de base de datos del mundo real, la carga de datos en Spanner hasta el paso 2 sería transaccional. Para mantener intactas las prácticas recomendadas de diseño, prefiero mantener normalizadas las tablas transaccionales, por lo que creo una tabla separada para las incorporaciones.

CREATE TABLE apparels_embeddings (id string(100), embedding ARRAY<FLOAT64>)
PRIMARY KEY (id);

INSERT INTO apparels_embeddings(id, embeddings) 
SELECT CAST(id as string), embeddings.values
FROM ML.PREDICT(
  MODEL text_embeddings,
  (SELECT id, content from apparels)
) ;

Ahora que el contenido masivo y los embeddings están listos, creemos un índice y un extremo de la Búsqueda de vectores para almacenar los embeddings que ayudarán a realizar la búsqueda de vectores.

4. Trabajo del flujo de trabajo: Exportación de datos de Spanner a la Búsqueda de vectores

Crea un bucket de Cloud Storage

Esto es necesario para almacenar las incorporaciones de Spanner en un bucket de GCS en un formato JSON que la Búsqueda vectorial espera como entrada. Crea un bucket en la misma región que tus datos en Spanner. Si es necesario, crea una carpeta dentro, pero, principalmente, crea un archivo vacío llamado empty.json.

Configura Cloud Workflow

Para configurar una exportación por lotes de Spanner a un índice de la Búsqueda de vectores de Vertex AI, haz lo siguiente:

Crea un índice vacío:

Asegúrate de que el índice de Vector Search esté en la misma región que tu bucket de Cloud Storage y los datos. Sigue los 11 pasos de instrucciones que se encuentran en la pestaña de la consola en la sección Crea un índice para la actualización por lotes de la página de administración de índices. En la carpeta que se pasa a contentsDeltaUri, crea un archivo vacío llamado empty.json, ya que no podrás crear un índice sin este archivo. Esto crea un índice vacío.

Si ya tienes un índice, puedes omitir este paso. El flujo de trabajo reemplazará tu índice.

Nota: No puedes implementar un índice vacío en un extremo. Por lo tanto, pospondremos el paso de implementarlo en un extremo para más adelante, después de exportar los datos vectoriales a Cloud Storage.

Clona este repositorio de Git: Hay varias formas de clonar un repositorio de Git. Una de ellas es ejecutar el siguiente comando con la CLI de GitHub. Ejecuta los siguientes 2 comandos desde la terminal de Cloud Shell:

gh repo clone cloudspannerecosystem/spanner-ai

cd spanner-ai/vertex-vector-search/workflows

Esta carpeta contiene dos archivos

batch-export.yaml: Esta es la definición del flujo de trabajo.
sample-batch-input.json: Esta es una muestra de los parámetros de entrada del flujo de trabajo.

Configura input.json desde el archivo de muestra: Primero, copia el archivo JSON de muestra.

cp sample-batch-input.json input.json

Luego, edita input.json con los detalles de tu proyecto. En este caso, tu JSON debería verse así:

{
  "project_id": "<<YOUR_PROJECT>>",
  "location": "<<us-central1>>",
  "dataflow": {
    "temp_location": "gs://<<YOUR_BUCKET>>/<<FOLDER_IF_ANY>>/workflow_temp"
  },
  "gcs": {
    "output_folder": "gs://<<YOUR_BUCKET>>/<<FOLDER_IF_ANY>>/workflow_output"
  },
  "spanner": {
    "instance_id": "spanner-vertex",
    "database_id": "spanner-vertex-embeddings",
    "table_name": "apparels_embeddings",
    "columns_to_export": "embedding,id"
  },
  "vertex": {
    "vector_search_index_id": "<<YOUR_INDEX_ID>>"
  }
}

Configurar permisos

Para los entornos de producción, recomendamos crear una cuenta de servicio nueva y otorgarle uno o más roles de IAM que contengan los permisos mínimos necesarios para administrar el servicio. Se necesitan los siguientes roles para configurar el flujo de trabajo para exportar datos de Spanner (incorporaciones) al índice de Vector Search:

Cuenta de servicio de Cloud Workflows:

De forma predeterminada, se usa la cuenta de servicio predeterminada de Compute Engine.

Si usas una cuenta de servicio configurada manualmente, debes incluir los siguientes roles:

Para activar un trabajo de Dataflow: Administrador de Dataflow, trabajador de Dataflow.

Para suplantar una cuenta de servicio de trabajador de Dataflow: Usuario de cuenta de servicio.

Para escribir registros: Escritor de registros.

Para activar la recompilación de Vertex AI Vector Search: Usuario de Vertex AI.

Cuenta de servicio de trabajador de Dataflow:

Si usas una cuenta de servicio configurada manualmente, debes incluir los siguientes roles:

Para administrar Dataflow: Administrador de Dataflow y Trabajador de Dataflow. Para leer datos de Spanner, usa Cloud Spanner Database Reader. Acceso de escritura al registro de contenedores de GCS seleccionado: Propietario del bucket de almacenamiento de GCS.

Implementa el flujo de trabajo de Cloud

Implementa el archivo YAML de flujo de trabajo en tu proyecto de Google Cloud. Puedes configurar la región o la ubicación en la que se ejecutará el flujo de trabajo.

gcloud workflows deploy vector-export-workflow --source=batch-export.yaml --location="us-central1" [--service account=<service_account>]

or 

gcloud workflows deploy vector-export-workflow --source=batch-export.yaml --location="us-central1"

Ahora, el flujo de trabajo debería estar visible en la página Workflows en la consola de Google Cloud.

Nota: También puedes crear e implementar el flujo de trabajo desde la consola de Google Cloud. Sigue las indicaciones en la consola de Cloud. Para la definición del flujo de trabajo, copia y pega el contenido de batch-export.yaml.

Una vez que se complete este paso, ejecuta el flujo de trabajo para que comience la exportación de datos.

Ejecuta el flujo de trabajo de Cloud

Ejecuta el siguiente comando para ejecutar el flujo de trabajo:

gcloud workflows execute vector-export-workflow --data="$(cat input.json)"

La ejecución debería aparecer en la pestaña Ejecuciones de Workflows. Esto debería cargar tus datos en la base de datos de Vector Search y, luego, indexarlos.

Nota: También puedes ejecutar desde la consola con el botón Ejecutar. Sigue las instrucciones y, para la entrada, copia y pega el contenido de tu archivo input.json personalizado.

5. Implementa el índice de Vector Search

Implementa el índice en un extremo

Puedes seguir los pasos que se indican a continuación para implementar el índice:

En la página Índices de Vector Search, deberías ver un botón DEPLOY junto al índice que acabas de crear en el paso 2 de la sección anterior. También puedes navegar a la página de información del índice y hacer clic en el botón IMPLEMENTAR EN EL EXTREMO.
Proporciona la información necesaria y, luego, implementa el índice en un extremo.

Como alternativa, puedes consultar este notebook para implementarlo en un extremo (ve a la parte de implementación del notebook). Una vez que se implemente, anota el ID del índice implementado y la URL del extremo.

6. Frontend: Datos del usuario para la Búsqueda de Vectores

Creemos una aplicación de Python simple con una UX potenciada por Gradio para probar rápidamente nuestra implementación. Puedes consultar la implementación aquí para implementar esta app de demostración en tu propio notebook de Colab.

Usaremos el SDK de aiplatform de Python para llamar a la API de Embeddings y también para invocar el extremo del índice de Vector Search.

# [START aiplatform_sdk_embedding]
!pip install google-cloud-aiplatform==1.35.0 --upgrade --quiet --user


import vertexai
vertexai.init(project=PROJECT_ID, location="us-central1")


from vertexai.language_models import TextEmbeddingModel


import sys
if "google.colab" in sys.modules:
    # Define project information
    PROJECT_ID = " "  # Your project id
    LOCATION = " "  # Your location 


    # Authenticate user to Google Cloud
    from google.colab import auth
    auth.authenticate_user()

Usaremos Gradio para mostrar la aplicación de IA que estamos creando de forma rápida y sencilla con una interfaz de usuario. Reinicia el entorno de ejecución antes de implementar este paso.

!pip install gradio
import gradio as gr

Desde la app web, cuando el usuario ingresa datos, invoca la API de Embeddings. Usaremos el modelo de embedding de texto: textembedding-gecko@latest

El siguiente método invoca el modelo de Text Embedding y devuelve los embeddings de vectores para el texto que ingresó el usuario:

def text_embedding(content) -> list:
    """Text embedding with a Large Language Model."""
    model = TextEmbeddingModel.from_pretrained("textembedding-gecko@latest")
    embeddings = model.get_embeddings(content)
    for embedding in embeddings:
        vector = embedding.values
        #print(f"Length of Embedding Vector: {len(vector)}")
    return vector

Pruébalo

text_embedding("red shorts for girls")

Deberías ver un resultado similar al siguiente (ten en cuenta que la imagen está recortada en altura, por lo que no puedes ver toda la respuesta del vector):

Declara el ID del índice implementado y el ID del extremo

from google.cloud import aiplatform
DEPLOYED_INDEX_ID = "spanner_vector1_1702366982123"

#Vector Search Endpoint
index_endpoint = aiplatform.MatchingEngineIndexEndpoint('projects/273845608377/locations/us-central1/indexEndpoints/2021628049526620160')

Define el método de Vector Search para llamar al extremo del índice y mostrar el resultado con las 10 coincidencias más cercanas para la respuesta de la incorporación correspondiente al texto de entrada del usuario.

En la siguiente definición del método para la búsqueda vectorial, ten en cuenta que se invoca el método find_neighbors para identificar los 10 vectores más cercanos.

def vector_search(content) -> list:
  result = text_embedding(content)
  #call_vector_search_api(content)
  index_endpoint = aiplatform.MatchingEngineIndexEndpoint('projects/273845608377/locations/us-central1/indexEndpoints/2021628049526620160')
  # run query
  response = index_endpoint.find_neighbors(
      deployed_index_id = DEPLOYED_INDEX_ID,
      queries = [result],
      num_neighbors = 10
  )
  out = []
  # show the results
  for idx, neighbor in enumerate(response[0]):
      print(f"{neighbor.distance:.2f} {spanner_read_data(neighbor.id)}")
      out.append(f"{spanner_read_data(neighbor.id)}")
  return out

También notarás la llamada al método spanner_read_data. Lo veremos en el siguiente paso.

Define la implementación del método de lectura de datos de Spanner que invoca el método execute_sql para extraer las imágenes correspondientes a los IDs de los vectores de vecinos más cercanos que se devolvieron en el último paso.

!pip install google-cloud-spanner==3.36.0


from google.cloud import spanner


instance_id = "spanner-vertex"
database_id = "spanner-vertex-embeddings"
projectId = PROJECT_ID
client = spanner.Client()
client.project = projectId
instance = client.instance(instance_id)
database = instance.database(database_id)
def spanner_read_data(id):
    query = "SELECT uri FROM apparels where id = " + id
    outputs = []
    with database.snapshot() as snapshot:
        results = snapshot.execute_sql(query)


        for row in results:
            #print(row)
            #output = "ID: {}, CONTENT: {}, URI: {}".format(*row)
            output = "{}".format(*row)
            outputs.append(output)


    return "\n".join(outputs)

Debería devolver las URLs de las imágenes correspondientes a los vectores elegidos.

Por último, unamos las piezas en una interfaz de usuario y activemos el proceso de Vector Search.

from PIL import Image


def call_search(query):
  response = vector_search(query)
  return response


input_text = gr.Textbox(label="Enter your query. Examples: Girls Tops White Casual, Green t-shirt girls, jeans shorts, denim skirt etc.")
output_texts = [gr.Image(label="") for i in range(10)]
demo = gr.Interface(fn=call_search, inputs=input_text, outputs=output_texts, live=True)
resp = demo.launch(share = True)

Deberías ver el resultado que se muestra a continuación:

Imagen: Vínculo

Mira el video del resultado aquí.

7. Limpia

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta publicación:

En la consola de Google Cloud, ve a la página Administrar recursos.
En la lista de proyectos, elige el proyecto que quieres borrar y haz clic en Borrar.
En el diálogo, escribe el ID del proyecto y, luego, haz clic en Cerrar para borrarlo.
Si no quieres borrar el proyecto, borra la instancia de Spanner. Para ello, navega a la instancia que acabas de crear para este proyecto y haz clic en el botón BORRAR INSTANCIA en la esquina superior derecha de la página de descripción general de la instancia.
También puedes navegar al índice de Vector Search, anular la implementación del extremo y el índice, y borrar el índice.

8. Conclusión

¡Felicitaciones! Completaste correctamente la implementación de Spanner - Vertex Vector Search con los siguientes pasos:

Crear una fuente de datos y una incorporación de Spanner para las aplicaciones que se obtienen de la base de datos de Spanner
Se está creando el índice de la base de datos de Vector Search.
Integrar datos de vectores de Spanner en la Búsqueda de vectores con trabajos de Dataflow y Workflow
Implementar el índice en un extremo
Por último, invoca la Búsqueda vectorial en la entrada del usuario en una implementación potenciada por Python del SDK de Vertex AI.

No dudes en ampliar la implementación a tu propio caso de uso o improvisar el caso de uso actual con nuevas funciones. Obtén más información sobre las capacidades de aprendizaje automático de Spanner aquí.