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/5] 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/5] - 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 From 0d875ad631bffb8ad57088d4066d4fcba5f5336b Mon Sep 17 00:00:00 2001 From: David del Real Sifuentes Date: Thu, 6 Aug 2026 16:56:31 +0000 Subject: [PATCH 3/5] - Migrated model_tuning_example to genai/embeddings. - created tests for migrated samples. --- genai/embeddings/code_retrieval_example.py | 46 ++++++++++---------- genai/embeddings/model_tuning_example.py | 32 +++++++------- genai/embeddings/test_embeddings_examples.py | 11 +++++ 3 files changed, 50 insertions(+), 39 deletions(-) diff --git a/genai/embeddings/code_retrieval_example.py b/genai/embeddings/code_retrieval_example.py index bbdd2d03f18..a56ecd34346 100644 --- a/genai/embeddings/code_retrieval_example.py +++ b/genai/embeddings/code_retrieval_example.py @@ -18,7 +18,6 @@ from google import genai - # TODO (Developer) set the following environment variables. PROJECT_ID = os.getenv("PROJECT_ID") LOCATION_ID = os.getenv("LOCATION_ID", "us-central1") @@ -29,53 +28,52 @@ 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 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[genai.types.EmbedContentConfig, genai.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 + 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 for line in SOURCE_CODE: - config = genai.types.EmbedContentConfig( - task_type=RETRIEVAL_DOCUMENT - ) + config = genai.types.EmbedContentConfig(task_type=RETRIEVAL_DOCUMENT) index_response = client.models.embed_content( - model=MODEL_NAME, - contents=line, - config=config + model=MODEL_NAME, contents=line, config=config ) - print(f"Task: {RETRIEVAL_DOCUMENT} | " - f"Vector length: {len(index_response.embeddings)} | " - f"Preview: {index_response.embeddings[:3]}...") + print( + f"Task: {RETRIEVAL_DOCUMENT} | " + f"Vector length: {len(index_response.embeddings)} | " + f"Preview: {index_response.embeddings[:3]}..." + ) # Embed Search Prompts for line in QUERY_LINES: - config = genai.types.EmbedContentConfig( - task_type=CODE_RETRIEVAL_QUERY - ) + config = genai.types.EmbedContentConfig(task_type=CODE_RETRIEVAL_QUERY) query_response = client.models.embed_content( - model=MODEL_NAME, - contents=line, - config=config + model=MODEL_NAME, contents=line, config=config ) - print(f"Task: {CODE_RETRIEVAL_QUERY} | " - f"Vector length: {len(query_response.embeddings)} | " - f"Preview: {query_response.embeddings[:3]}...") + print( + f"Task: {CODE_RETRIEVAL_QUERY} | " + f"Vector length: {len(query_response.embeddings)} | " + f"Preview: {query_response.embeddings[:3]}..." + ) return index_response, query_response + # [END aiplatform_genai_embedding_code_retrieval] diff --git a/genai/embeddings/model_tuning_example.py b/genai/embeddings/model_tuning_example.py index 1d678ba1cbc..68549915ce3 100644 --- a/genai/embeddings/model_tuning_example.py +++ b/genai/embeddings/model_tuning_example.py @@ -17,38 +17,39 @@ 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") +MODEL_NAME = os.getenv("MODEL_NAME", "text-embedding-004") # 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" +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" -) +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. - """ + """Tune an embedding model using the specified parameters.""" - aiplatform.init( - project=PROJECT_ID, - location=LOCATION_ID - ) + aiplatform.init(project=PROJECT_ID, location=LOCATION_ID) # Configure parameters expected by the embedding tuning pipeline template pipeline_parameters = { @@ -77,4 +78,5 @@ def tune_embedding_model() -> aiplatform.PipelineJob: return pipeline_job + # [END aiplatform_genai_embedding_model_tuning] diff --git a/genai/embeddings/test_embeddings_examples.py b/genai/embeddings/test_embeddings_examples.py index 5b8d94f07c3..b3a50ef105a 100644 --- a/genai/embeddings/test_embeddings_examples.py +++ b/genai/embeddings/test_embeddings_examples.py @@ -18,7 +18,9 @@ import os +import code_retrieval_example import embeddings_docretrieval_with_txt +import model_tuning_example os.environ["GOOGLE_GENAI_USE_ENTERPRISE"] = "True" os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1" @@ -29,3 +31,12 @@ def test_embeddings_docretrieval_with_txt() -> None: response = embeddings_docretrieval_with_txt.embed_content() assert response + + +def test_code_retrieva_example() -> None: + response = code_retrieval_example.embed_test() + assert response + +def test_model_tuning_example() -> None: + response = model_tuning_example.tune_embedding_model() + assert response From f927421a6db35ae25325413be51d05ebca7b9f28 Mon Sep 17 00:00:00 2001 From: David del Real Sifuentes Date: Thu, 6 Aug 2026 17:17:48 +0000 Subject: [PATCH 4/5] - Applied Gemini review code suggestions. --- genai/embeddings/code_retrieval_example.py | 3 +-- genai/embeddings/model_tuning_example.py | 3 --- genai/embeddings/test_embeddings_examples.py | 2 +- 3 files changed, 2 insertions(+), 6 deletions(-) diff --git a/genai/embeddings/code_retrieval_example.py b/genai/embeddings/code_retrieval_example.py index a56ecd34346..7446a8a6ed3 100644 --- a/genai/embeddings/code_retrieval_example.py +++ b/genai/embeddings/code_retrieval_example.py @@ -22,7 +22,6 @@ 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" @@ -35,7 +34,7 @@ def embed_test() -> ( - tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentConfig] + tuple[genai.types.EmbedContentConfig, genai.types.EmbedContentResponse] ): """Generates embeddings for source code indexing and code search queries using the Gemini API. diff --git a/genai/embeddings/model_tuning_example.py b/genai/embeddings/model_tuning_example.py index 68549915ce3..12bdb7c8a12 100644 --- a/genai/embeddings/model_tuning_example.py +++ b/genai/embeddings/model_tuning_example.py @@ -24,9 +24,6 @@ # 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" ) diff --git a/genai/embeddings/test_embeddings_examples.py b/genai/embeddings/test_embeddings_examples.py index b3a50ef105a..a6afaa00701 100644 --- a/genai/embeddings/test_embeddings_examples.py +++ b/genai/embeddings/test_embeddings_examples.py @@ -33,7 +33,7 @@ def test_embeddings_docretrieval_with_txt() -> None: assert response -def test_code_retrieva_example() -> None: +def test_code_retrieval_example() -> None: response = code_retrieval_example.embed_test() assert response From c4f0c6510f9afc0ba5ca0c5de8397079e3a27d86 Mon Sep 17 00:00:00 2001 From: David del Real Sifuentes Date: Thu, 6 Aug 2026 19:03:02 +0000 Subject: [PATCH 5/5] Removed unused attribute. --- genai/embeddings/model_tuning_example.py | 1 - 1 file changed, 1 deletion(-) diff --git a/genai/embeddings/model_tuning_example.py b/genai/embeddings/model_tuning_example.py index 12bdb7c8a12..c6d7c268ef2 100644 --- a/genai/embeddings/model_tuning_example.py +++ b/genai/embeddings/model_tuning_example.py @@ -55,7 +55,6 @@ def tune_embedding_model() -> aiplatform.PipelineJob: "queries_path": QUERIES_PATH, "train_label_path": TRAIN_LABEL_PATH, "test_label_path": TEST_LABEL_PATH, - "batch_size": BATCH_SIZE, "accelerator_type": ACCELERATOR_TYPE, }