From 3348d292815ad70946f10cdd2251a86db7568bd6 Mon Sep 17 00:00:00 2001 From: David del Real Sifuentes Date: Tue, 4 Aug 2026 23:05:15 +0000 Subject: [PATCH 1/2] Migrate code retrieval example for embeddings to genai. --- genai/embeddings/code_retrieval_example.py | 80 ++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 genai/embeddings/code_retrieval_example.py diff --git a/genai/embeddings/code_retrieval_example.py b/genai/embeddings/code_retrieval_example.py new file mode 100644 index 00000000000..de74116669d --- /dev/null +++ b/genai/embeddings/code_retrieval_example.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START aiplatform_genai_embedding_code_retrieval] + +import os + +from google import genai +from google.genai import types + + +# TODO (Developer) set the following environment variables. +PROJECT_ID = os.getenv("PROJECT_ID") +LOCATION_ID = os.getenv("LOCATION_ID", "us-central1") +MODEL_NAME = os.getenv("MODEL_NAME", "gemini-embedding-001") +DIMENSIONALITY = 3072 + +QUERY_LINES = ["Retrieve a function that adds two numbers"] +CODE_RETRIEVAL_QUERY = "CODE_RETRIEVAL_QUERY" +RETRIEVAL_DOCUMENT = "RETRIEVAL_DOCUMENT" +SOURCE_CODE = [ + "def func(a, b): return a + b", + "def func(a, b): return a - b", + "def func(a, b): return (a ** 2 + b ** 2) ** 0.5", +] + + +def embed_test() -> tuple[types.EmbedContentConfig, types.EmbedContentConfig]: + client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION_ID) + + # Index Source Code + for line in SOURCE_CODE: + config = genai.types.EmbedContentConfig( + task_type=RETRIEVAL_DOCUMENT + ) + + index_response = client.models.embed_content( + model=MODEL_NAME, + contents=line, + config=config + ) + + print(f"Task: {RETRIEVAL_DOCUMENT} | " + f"Vector length: {len(index_response.embedding.values)} | " + f"Preview: {index_response.embedding.values[:3]}...") + + # Embed Search Prompts + for line in QUERY_LINES: + config = genai.types.EmbedContentConfig( + task_type=CODE_RETRIEVAL_QUERY + ) + + query_response = client.models.embed_content( + model=MODEL_NAME, + contents=line, + config=config + ) + + print(f"Task: {CODE_RETRIEVAL_QUERY} | " + f"Vector length: {len(query_response.embedding.values)} | " + f"Preview: {query_response.embedding.values[:3]}...") + + return index_response, query_response + + +if __name__ == "__main__": + embed_test() + +# [END aiplatform_genai_embedding_code_retrieval] From f848b0726409f4679268dfd135bc6fa0c5f1dad1 Mon Sep 17 00:00:00 2001 From: David del Real Sifuentes Date: Wed, 5 Aug 2026 23:18:10 +0000 Subject: [PATCH 2/2] - Migrated model_tuning_example sample to modern aiplatform implementation of embedded model tuning. - Added missing python doc. - Updated requirements files. --- genai/embeddings/code_retrieval_example.py | 21 +++--- genai/embeddings/model_tuning_example.py | 80 ++++++++++++++++++++++ genai/embeddings/requirements-test.txt | 4 +- genai/embeddings/requirements.txt | 3 +- 4 files changed, 95 insertions(+), 13 deletions(-) create mode 100644 genai/embeddings/model_tuning_example.py diff --git a/genai/embeddings/code_retrieval_example.py b/genai/embeddings/code_retrieval_example.py index de74116669d..bbdd2d03f18 100644 --- a/genai/embeddings/code_retrieval_example.py +++ b/genai/embeddings/code_retrieval_example.py @@ -17,7 +17,6 @@ import os from google import genai -from google.genai import types # TODO (Developer) set the following environment variables. @@ -36,7 +35,13 @@ ] -def embed_test() -> tuple[types.EmbedContentConfig, types.EmbedContentConfig]: +def embed_test() -> tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentConfig]: + """Generates embeddings for source code indexing and code search queries using the Gemini API. + + Returns: + tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentConfig]: A tuple containing + the final source code indexing response and search query embedding response. + """ client = genai.Client(enterprise=True, project=PROJECT_ID, location=LOCATION_ID) # Index Source Code @@ -52,8 +57,8 @@ def embed_test() -> tuple[types.EmbedContentConfig, types.EmbedContentConfig]: ) print(f"Task: {RETRIEVAL_DOCUMENT} | " - f"Vector length: {len(index_response.embedding.values)} | " - f"Preview: {index_response.embedding.values[:3]}...") + f"Vector length: {len(index_response.embeddings)} | " + f"Preview: {index_response.embeddings[:3]}...") # Embed Search Prompts for line in QUERY_LINES: @@ -68,13 +73,9 @@ def embed_test() -> tuple[types.EmbedContentConfig, types.EmbedContentConfig]: ) print(f"Task: {CODE_RETRIEVAL_QUERY} | " - f"Vector length: {len(query_response.embedding.values)} | " - f"Preview: {query_response.embedding.values[:3]}...") + f"Vector length: {len(query_response.embeddings)} | " + f"Preview: {query_response.embeddings[:3]}...") return index_response, query_response - -if __name__ == "__main__": - embed_test() - # [END aiplatform_genai_embedding_code_retrieval] diff --git a/genai/embeddings/model_tuning_example.py b/genai/embeddings/model_tuning_example.py new file mode 100644 index 00000000000..1d678ba1cbc --- /dev/null +++ b/genai/embeddings/model_tuning_example.py @@ -0,0 +1,80 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# [START aiplatform_genai_embedding_model_tuning] +import os + +from google.cloud import aiplatform + + +# TODO (Developer) set the following environment variables. +PROJECT_ID = os.getenv("PROJECT_ID") +LOCATION_ID = os.getenv("LOCATION_ID", "us-central1") +MODEL_NAME = os.getenv("MODEL_NAME", "text-embedding-005") +# A storage bucket: gs://your-bucket-name/embedding-tuning-output +OUTPUT_URI = os.getenv("OUTPUT_DIR") + +API_ENDPOINT = f"{LOCATION_ID}-aiplatform.googleapis.com" +BATCH_SIZE = 128 +LEARNING_RATE_MULTIPLIER = 1.0 +TRAIN_LABEL_PATH = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/train.tsv" +TEST_LABEL_PATH = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/test.tsv" +CORPUS_PATH = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/corpus.jsonl" +QUERIES_PATH = "gs://cloud-samples-data/ai-platform/embedding/goog-10k-2024/r11/queries.jsonl" + +ACCELERATOR_TYPE = "NVIDIA_L4" + +# Official Google Cloud KFP pipeline template URI for text embedding model tuning +EMBEDDING_TUNING_PIPELINE_URI = ( + "https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.3" +) + + +def tune_embedding_model() -> aiplatform.PipelineJob: + """Tune an embedding model using the specified parameters. + """ + + aiplatform.init( + project=PROJECT_ID, + location=LOCATION_ID + ) + + # Configure parameters expected by the embedding tuning pipeline template + pipeline_parameters = { + "base_model_version_id": MODEL_NAME, + "corpus_path": CORPUS_PATH, + "queries_path": QUERIES_PATH, + "train_label_path": TRAIN_LABEL_PATH, + "test_label_path": TEST_LABEL_PATH, + "batch_size": BATCH_SIZE, + "accelerator_type": ACCELERATOR_TYPE, + } + + # Instantiate the Vertex AI Pipeline job + pipeline_job = aiplatform.PipelineJob( + display_name="tune-text-embedding-model-job", + template_path=EMBEDDING_TUNING_PIPELINE_URI, + pipeline_root=OUTPUT_URI, + parameter_values=pipeline_parameters, + project=PROJECT_ID, + location=LOCATION_ID, + ) + + pipeline_job.submit() + + print(f"Pipeline submitted successfully: {pipeline_job.resource_name}") + + return pipeline_job + +# [END aiplatform_genai_embedding_model_tuning] diff --git a/genai/embeddings/requirements-test.txt b/genai/embeddings/requirements-test.txt index 22a9617b8e8..966a12401f2 100644 --- a/genai/embeddings/requirements-test.txt +++ b/genai/embeddings/requirements-test.txt @@ -1,2 +1,2 @@ -google-api-core==2.24.0 -pytest==9.0.3; python_version >= "3.10" +google-api-core==2.33.0 +pytest==9.1.1 diff --git a/genai/embeddings/requirements.txt b/genai/embeddings/requirements.txt index 1efe7b29dbc..88bf66e0d7d 100644 --- a/genai/embeddings/requirements.txt +++ b/genai/embeddings/requirements.txt @@ -1 +1,2 @@ -google-genai==1.42.0 +google-genai==2.16.0 +google-cloud-aiplatform[pipelines]==1.163.0