feat: Text-Embedding Model Tuning demo. (#11368)

* feat: Text-Embedding Model Tuning demo.

* feat: Text-Embedding Model Tuning demo.

* feat: Text-Embedding Model Tuning demo.

* feat: Text-Embedding Model Tuning demo.

generative_ai/embedding_model_tuning.py

@@ -0,0 +1,63 @@
+# Copyright 2024 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_sdk_embedding]
+import re
+
+from google.cloud import aiplatform
+from google.cloud.aiplatform import initializer as aiplatform_init
+from google.cloud.aiplatform import pipeline_jobs
+
+
+def tune_embedding_model(
+    api_endpoint: str,
+    project: str,
+    output_dir: str,
+    pipeline_job_display_name: str = "embedding-customization-pipeline-sample",
+    base_model_version_id: str = "textembedding-gecko@003",
+    task_type: str = "DEFAULT",
+    queries_path: str = "gs://embedding-customization-pipeline/dataset/queries.jsonl",
+    corpus_path: str = "gs://embedding-customization-pipeline/dataset/corpus.jsonl",
+    train_label_path: str = "gs://embedding-customization-pipeline/dataset/train.tsv",
+    test_label_path: str = "gs://embedding-customization-pipeline/dataset/test.tsv",
+    batch_size: int = 50,
+    iterations: int = 300
+) -> pipeline_jobs.PipelineJob:
+    match = re.search(r"(.+)(-autopush|-staging)?-aiplatform.+", api_endpoint)
+    location = match.group(1) if match else "us-central1"
+    job = aiplatform.PipelineJob(
+        display_name=pipeline_job_display_name,
+        template_path="https://us-kfp.pkg.dev/ml-pipeline/llm-text-embedding/tune-text-embedding-model/v1.1.2",
+        pipeline_root=output_dir,
+        parameter_values=dict(
+            project=project,
+            location=location,
+            base_model_version_id=base_model_version_id,
+            task_type=task_type,
+            queries_path=queries_path,
+            corpus_path=corpus_path,
+            train_label_path=train_label_path,
+            test_label_path=test_label_path,
+            batch_size=batch_size,
+            iterations=iterations)
+    )
+    job.submit()
+    return job
+
+
+# [END aiplatform_sdk_embedding]
+if __name__ == "__main__":
+    tune_embedding_model(aiplatform_init.global_config.api_endpoint,
+                         aiplatform_init.global_config.project,
+                         aiplatform_init.global_config.staging_bucket)

generative_ai/embedding_model_tuning_test.py

@@ -0,0 +1,51 @@
+# Copyright 2024 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.
+
+import os
+
+import backoff
+
+from google.api_core.exceptions import FailedPrecondition
+import google.auth
+from google.cloud import aiplatform
+from google.cloud.aiplatform import initializer as aiplatform_init
+from google.cloud.aiplatform import pipeline_jobs
+
+
+import embedding_model_tuning
+
+
+@backoff.on_exception(backoff.expo, FailedPrecondition, max_time=300)
+def dispose(job: pipeline_jobs.PipelineJob) -> None:
+    if not job.done():
+        job.cancel()
+    job.delete()
+
+
+def test_tune_embedding_model() -> None:
+    credentials, _ = google.auth.default(  # Set explicit credentials with Oauth scopes.
+        scopes=["https://www.googleapis.com/auth/cloud-platform"])
+    aiplatform.init(
+        api_endpoint="us-central1-aiplatform.googleapis.com:443",
+        project=os.getenv("GOOGLE_CLOUD_PROJECT"),
+        staging_bucket="gs://ucaip-samples-us-central1/training_pipeline_output",
+        credentials=credentials)
+    job = embedding_model_tuning.tune_embedding_model(
+        aiplatform_init.global_config.api_endpoint,
+        aiplatform_init.global_config.project,
+        aiplatform_init.global_config.staging_bucket)
+    try:
+        assert job.state != "PIPELINE_STATE_FAILED"
+    finally:
+        dispose(job)