Crea una pipeline di vector embedding in tempo reale per AlloyDB con Dataflow

Questo documento mostra come creare una pipeline ETL (Extract, Transform, Load) di AlloyDB utilizzando Dataflow. Google Cloud Dataflow è un servizio completamente gestito Google Cloud per lo sviluppo e l'esecuzione di pipeline di trattamento dati.

Puoi utilizzare le istruzioni del documento, basate sul notebook Colab Vector Embedding Ingestion with Apache Beam and AlloyDB, che utilizza Python per creare la pipeline di importazione basic_ingestion_pipeline.py. Alcuni dei casi d'uso in cui puoi applicare le informazioni contenute in questo documento sono la ricerca semantica o la Retrieval Augmented Generation (RAG).

Queste istruzioni descrivono i seguenti componenti della pipeline Dataflow:

  • Configurazione di una connessione AlloyDB e Dataflow
  • Generazione di embedding in AlloyDB utilizzando il gestore VertexAITextEmbeddings di Apache Beam e il modello di text embedding della piattaforma Gemini Enterprise Agent
  • Creazione di una pipeline di inserimento flussi in Dataflow

Prima di iniziare

Prima di creare la pipeline Dataflow utilizzando Colab, completa i seguenti prerequisiti:

Configura l'istanza AlloyDB e i componenti della pipeline

Innanzitutto, configura la pipeline per la connessione a un'istanza AlloyDB. Questa configurazione include la definizione dell' Google Cloud ID progetto, dell'URI dell'istanza AlloyDB, dell'utente e della password per la connessione utilizzando il connettore di linguaggio AlloyDB. Per saperne di più sulla configurazione della connessione, consulta Configurazione del database.

I moduli Apache Beam specifici per la Retrieval Augmented Generation (RAG) forniscono classi per le seguenti attività:

  • Importazione di dati da AlloyDB
  • Generazione di embedding
  • Scrittura di questi vector embedding in AlloyDB

Importa le classi richieste nel codice della pipeline prima di creare la logica della pipeline. Per saperne di più sui componenti della pipeline, consulta Importare i componenti della pipeline.

Crea dati di esempio

Il notebook Colab Vector Embedding Ingestion with Apache Beam and AlloyDB fornisce dati di esempio products_data per l'esecuzione della pipeline. La pipeline utilizza questi dati di esempio come input, insieme al modello di embedding, per generare gli embedding.

Per saperne di più, consulta Creare dati di esempio.

Crea una tabella per archiviare gli embedding

La pipeline archivia gli embedding generati nella tabella default_dataflow_product_embeddings. Per saperne di più sulla creazione dello schema della tabella, consulta Creare una tabella con lo schema predefinito.

(Facoltativo) Prepara i dati per l'importazione degli embedding

A seconda del set di dati, puoi suddividere i dati in metadati e testo che il modello di embedding deve convertire in embedding. Le classi MLTransform() e VectorDatabaseWriteTransform() elaborano i dati di input in una dimensione supportata dal modello di embedding. Includi i metadati e formatta i dati di input in base alle specifiche del modello di embedding che stai utilizzando.

Per saperne di più sulla preparazione dei dati, consulta Mappare i dati dei prodotti in blocchi.

Configura il gestore di embedding per generare gli embedding

La classe VertexAITextEmbeddings() definisce il modello di text embedding che crea i vector embedding. Questo modello di embedding converte i dati suddivisi in blocchi in embedding.

Per saperne di più, consulta Configurare il gestore di embedding.

Puoi anche utilizzare un modello preaddestrato creato con il framework Huggingface SentenceTransformers per generare vector embedding. Per saperne di più, consulta Generare embedding con HuggingFace.

Crea una pipeline di importazione

La pipeline basic_ingestion_pipeline.py, fornita nel notebook Colab Vector Embedding Ingestion with Apache Beam and AlloyDB, incorpora le configurazioni delle sezioni precedenti, tra cui la configurazione di AlloyDB, il caricamento dei dati in AlloyDB, la suddivisione facoltativa dei dati in blocchi e la configurazione del gestore di embedding.

La pipeline di importazione esegue le seguenti operazioni:

  • Crea tabelle di dati dei prodotti
  • Converte i dati in blocchi
  • Genera embedding
  • Scrive gli embedding convertiti nella tabella products_data in AlloyDB

Puoi eseguire questa pipeline utilizzando un runner locale diretto o un runner basato sul cloud come Dataflow.

Per saperne di più sulla creazione della pipeline di importazione, consulta Salvare la pipeline in un file Python.

Esegui la pipeline Dataflow

Puoi eseguire una pipeline Dataflow dalla riga di comando. Trasmetti le credenziali, come l'ID progetto, i dettagli di connessione AlloyDB, la località del bucket Cloud Storage, i dettagli dell'ambiente di esecuzione, le informazioni sulla rete e il nome della pipeline di importazione (basic_ingestion_pipeline.py).

Nel notebook Colab Vector Embedding Ingestion with Apache Beam and AlloyDB, l'istanza AlloyDB per PostgreSQL e i job Dataflow vengono eseguiti nella stessa rete VPC e nella stessa subnet.

Per saperne di più sull'esecuzione di una pipeline in Dataflow, consulta Eseguire la pipeline su Dataflow.

Nella Google Cloud console, nella dashboard Dataflow, puoi visualizzare grafici di esecuzione, log e metriche durante l'esecuzione della pipeline.

(Facoltativo) Esegui la pipeline Dataflow di inserimento flussi

Per i dati che dovrebbero cambiare spesso, come le ricerche di somiglianza o i motori di suggerimenti, valuta la possibilità di creare una pipeline di inserimento flussi utilizzando Dataflow e Pub/Sub.

Anziché elaborare un batch di dati, questa pipeline legge continuamente i messaggi in entrata da un argomento Pub/Sub, li converte in blocchi, genera embedding utilizzando un modello specificato (come Hugging Face o Agent Platform) e aggiorna la tabella AlloyDB.

Per saperne di più, consulta Aggiornamenti degli embedding di inserimento flussi da Pub/Sub.

Verifica i vector embedding in AlloyDB per PostgreSQL

Dopo l'esecuzione della pipeline, verifica che abbia scritto gli embedding nel database AlloyDB.

Per saperne di più, consulta Verificare gli embedding scritti.

Passaggi successivi