diff --git a/alloydb/notebooks/embeddings_batch_processing.ipynb b/alloydb/notebooks/embeddings_batch_processing.ipynb index 862656f1c7..e5118f01b0 100644 --- a/alloydb/notebooks/embeddings_batch_processing.ipynb +++ b/alloydb/notebooks/embeddings_batch_processing.ipynb @@ -77,7 +77,8 @@ " google-cloud-alloydb-connector[asyncpg]==1.4.0 \\\n", " sqlalchemy==2.0.36 \\\n", " pandas==2.2.3 \\\n", - " vertexai==1.70.0 \\\n", + " google-cloud-aiplatform==1.165.1 \\\n", + " google-genai==2.19.0 \\\n", " asyncio==3.4.3 \\\n", " greenlet==3.1.1 \\\n", " --quiet" @@ -792,119 +793,133 @@ }, "outputs": [], "source": [ - "from google.api_core.exceptions import ResourceExhausted\n", - "from typing import Union\n", - "from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel\n", + "from typing import Any, AsyncIterator, List, Optional, Union\n", "\n", + "from google import genai\n", + "from google.genai import types\n", "\n", "async def embed_text(\n", " batch_data: List[dict[str, Any]],\n", - " model: TextEmbeddingModel,\n", " cols_to_embed: List[str],\n", + " client: genai.Client,\n", + " model_name: str = \"text-embedding-004\",\n", " task_type: str = \"SEMANTIC_SIMILARITY\",\n", " retries: int = 100,\n", " delay: int = 30,\n", ") -> List[dict[str, Union[List[float], str]]]:\n", - " \"\"\"Embeds text data from a batch of records using a Vertex AI embedding model.\n", + " \"\"\"Embeds text data from a batch of records using the google-genai SDK.\n", "\n", " Args:\n", - " batch_data: A data batch containing records with text data to embed.\n", - " model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n", - " cols_to_embed: A list of column names containing the data to be embedded.\n", - " task_type: The task type for the embedding model. Defaults to\n", - " \"SEMANTIC_SIMILARITY\".\n", - " Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n", - " retries: The maximum number of times to retry embedding generation in case\n", - " of errors. Defaults to 100.\n", - " delay: The delay in seconds between retries. Defaults to 30.\n", + " batch_data: A data batch containing records with text data to embed.\n", + " cols_to_embed: A list of column names containing the data to be embedded.\n", + " model_name: The embedding model ID.\n", + " task_type: The task type for the embedding model.\n", + " retries: The maximum number of retry attempts in case of errors.\n", + " delay: The delay in seconds between retries.\n", + " client: An optional genai.Client instance for connection reuse.\n", "\n", " Returns:\n", - " A list of records containing ids and embeddings.\n", - " Example:\n", - " [\n", - " {\n", - " 'id': 'id1',\n", - " 'col1_embedding': [1.0, 1.1, ...],\n", - " 'col2_embedding': [2.0, 2.1, ...],\n", - " ...\n", - " },\n", - " ...\n", - " ]\n", - " where col1 and col2 are columns containing data to be embedded.\n", + " A list of records containing IDs and mapped column embeddings.\n", " Raises:\n", - " Exception: Raises the encountered exception if all retries fail.\n", + " Exception: Raises the encountered exception if all retries fail.\n", " \"\"\"\n", " logger = logging.getLogger(\"embed_objects\")\n", " global total_char_count\n", "\n", - " # Place all of the embeddings into a single list\n", + " # Extract non-empty text strings to embed\n", " inputs = []\n", " for row in batch_data:\n", " for col in cols_to_embed:\n", " if col in row and row[col]:\n", - " inputs.append(TextEmbeddingInput(row[col], task_type))\n", + " inputs.append(str(row[col]))\n", "\n", - " # Retry loop\n", " for attempt in range(retries):\n", " try:\n", - " # Get embeddings for the text data\n", - " embeddings = await model.get_embeddings_async(inputs)\n", + " # Asynchronous API call using client_instance.aio\n", + " response = await client.aio.models.embed_content(\n", + " model=model_name,\n", + " contents=inputs,\n", + " config=types.EmbedContentConfig(\n", + " task_type=task_type,\n", + " ),\n", + " )\n", + "\n", + " # Track character metrics\n", + " total_char_count += sum(len(text) for text in inputs)\n", "\n", - " # Increase total char count\n", - " total_char_count += sum([len(input.text) for input in inputs])\n", + " # Map response embeddings back to dataset record structure\n", + " embeddings_list = (\n", + " response.embeddings\n", + " if response.embeddings\n", + " else [response.embedding]\n", + " )\n", + " embedding_iter = iter(embeddings_list)\n", "\n", - " # group the results together by id\n", - " embedding_iter = iter(embeddings)\n", " results = []\n", " for row in batch_data:\n", - " r = {\"id\": row[\"id\"]}\n", + " r = {\"id\": row.get(\"id\")}\n", " for col in cols_to_embed:\n", " if col in row and row[col]:\n", - " r[f\"{col}_embedding\"] = str(next(embedding_iter).values)\n", + " embedding_obj = next(embedding_iter)\n", + " r[f\"{col}_embedding\"] = str(embedding_obj.values)\n", " else:\n", " r[f\"{col}_embedding\"] = None\n", " results.append(r)\n", + "\n", " return results\n", "\n", " except Exception as e:\n", - " if attempt < retries - 1: # Retry only if attempts are left\n", - " logger.warning(f\"Error: {e}. Retrying in {delay} seconds...\")\n", - " await asyncio.sleep(delay) # Wait before retrying\n", + " if attempt < retries - 1:\n", + " logger.warning(\n", + " f\"Error: {e}. Retrying in {delay} seconds (attempt {attempt + 1}/{retries})...\"\n", + " )\n", + " await asyncio.sleep(delay)\n", " else:\n", - " logger.error(f\"Failed to get embeddings for data: {batch_data} after {retries} attempts.\")\n", + " logger.error(\n", + " f\"Failed to get embeddings after {retries} attempts: {e}\"\n", + " )\n", + " raise e\n", + "\n", " return []\n", "\n", "\n", "async def embed_objects_concurrently(\n", " cols_to_embed: List[str],\n", " batch_data: AsyncIterator[List[dict[str, Any]]],\n", - " model: TextEmbeddingModel,\n", - " task_type: str,\n", + " client: genai.Client,\n", + " model_name: str = \"text-embedding-004\",\n", + " task_type: str = \"SEMANTIC_SIMILARITY\",\n", " max_concurrency: int = 5,\n", ") -> AsyncIterator[List[dict[str, Union[str, List[float]]]]]:\n", " \"\"\"Embeds text data concurrently from an asynchronous batch data generator.\n", "\n", " Args:\n", - " cols_to_embed: A list of column names containing the data to be embedded.\n", - " batch_data: A data batch containing records with text data to embed.\n", - " model: The Vertex AI `TextEmbeddingModel` to use for generating embeddings.\n", - " task_type: The task type for the embedding model.\n", - " Supported task types: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/task-types\n", - " max_concurrency: The maximum number of embedding tasks to run concurrently.\n", - " Defaults to 5.\n", + " cols_to_embed: A list of column names containing the data to be embedded.\n", + " batch_data: An async generator yielding data batches with records.\n", + " model_name: The embedding model ID.\n", + " task_type: The task type for the embedding model.\n", + " max_concurrency: The maximum number of concurrent tasks to execute.\n", + " client: An optional genai.Client instance for connection reuse.\n", + "\n", " Yields:\n", - " A list of records containing ids and embeddings.\n", + " A list of records containing IDs and mapped column embeddings.\n", " \"\"\"\n", " logger = logging.getLogger(\"embed_objects\")\n", "\n", - " # Keep track of pending tasks\n", " pending: set[asyncio.Task] = set()\n", " has_next = True\n", + "\n", " while pending or has_next:\n", " while len(pending) < max_concurrency and has_next:\n", " try:\n", " data = await batch_data.__anext__()\n", - " coro = embed_text(data, model, cols_to_embed, task_type)\n", + " coro = embed_text(\n", + " batch_data=data,\n", + " cols_to_embed=cols_to_embed,\n", + " model_name=model_name,\n", + " task_type=task_type,\n", + " client=client,\n", + " )\n", " pending.add(asyncio.ensure_future(coro))\n", " except StopAsyncIteration:\n", " has_next = False\n", @@ -915,7 +930,9 @@ " )\n", " for task in done:\n", " result = task.result()\n", - " logger.info(f\"Embedding task completed: Processed {len(result)} rows.\")\n", + " logger.info(\n", + " f\"Embedding task completed: Processed {len(result)} rows.\"\n", + " )\n", " yield result" ] }, @@ -1045,9 +1062,9 @@ }, "outputs": [], "source": [ - "import vertexai\n", "import time\n", - "from vertexai.language_models import TextEmbeddingModel\n", + "\n", + "from google import genai\n", "\n", "### Define variables ###\n", "\n", @@ -1080,40 +1097,40 @@ "):\n", " \"\"\"Orchestrates the end-to-end workflow for generating and storing embeddings.\n", "\n", - " The workflow includes the following major steps:\n", - "\n", - " 1. Data Retrieval: Fetches data from the database that requires embedding.\n", - " 2. Batching: Divides the data into batches for optimized processing.\n", - " 3. Embedding Generation: Generates embeddings concurrently for the batched\n", - " data using the Vertex AI model.\n", - " 4. Database Update: Updates the database concurrently with the generated\n", - " embeddings.\n", + " Workflow Steps:\n", + " 1. Connection Initialization: Initializes DB pool.\n", + " 2. Client Initialization: Instantiates shared GenAI client for workflow lifetime.\n", + " 3. Data Retrieval & Batching: Fetches and chunks source data into async streams.\n", + " 4. Embedding Generation: Uses google-genai client asynchronously and concurrently.\n", + " 5. Database Update: Stores output embeddings concurrently in the database.\n", "\n", " Args:\n", - " pool_size: The size of the database connection pool. Defaults to 10.\n", - " embed_data_concurrency: The maximum number of concurrent tasks for generating embeddings.\n", - " Defaults to 20.\n", - " batch_update_concurrency: The maximum number of concurrent tasks for updating the database.\n", - " Defaults to 10.\n", + " pool_size: The size of the database connection pool.\n", + " embed_data_concurrency: Max concurrent tasks for generating embeddings.\n", + " batch_update_concurrency: Max concurrent tasks for database updates.\n", " \"\"\"\n", - " # Set up connections to the database\n", + " # Set up database connection pool\n", " pool = await init_connection_pool(connector, database_name, pool_size=pool_size)\n", "\n", - " # Initialise VertexAI and the model to be used to generate embeddings\n", - " vertexai.init(project=project_id, location=region)\n", - " model = TextEmbeddingModel.from_pretrained(model_name)\n", + " # Initialize single GenAI client instance for the top-level workflow lifetime\n", + " client = genai.Client(vertexai=True, project=project_id, location=region)\n", "\n", " start_time = time.monotonic()\n", "\n", - " # Fetch source data from the database\n", + " # Fetch source data from database\n", " source_data = get_source_data(pool, cols_to_embed)\n", "\n", - " # Divide the source data into batches for efficient processing\n", + " # Divide source data into asynchronous batches\n", " batch_data = batch_source_data(source_data, cols_to_embed)\n", "\n", - " # Generate embeddings for the batched data concurrently\n", + " # Generate embeddings concurrently using the GenAI SDK\n", " embeddings_data = embed_objects_concurrently(\n", - " cols_to_embed, batch_data, model, task, max_concurrency=embed_data_concurrency\n", + " cols_to_embed=cols_to_embed,\n", + " batch_data=batch_data,\n", + " model_name=model_name,\n", + " task_type=task,\n", + " max_concurrency=embed_data_concurrency,\n", + " client=client,\n", " )\n", "\n", " # Update the database with the generated embeddings concurrently\n", @@ -1124,7 +1141,7 @@ " end_time = time.monotonic()\n", " elapsed_time = end_time - start_time\n", "\n", - " # Release database connections and close the connector\n", + " # Release database resources\n", " await pool.dispose()\n", " await connector.close()\n", "\n",