refactor: (GenAI) Reorganized Embeddings Samples (Group A) (GoogleCloudPlatform#12575)

* New reorganized files

* Files to delete

* Update region tags

* Fix lint

* Changed "Example output" -> "Example response"

Author: Thoughtseize1

generative_ai/embedding.py

@@ -11,6 +11,7 @@
 # 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.
+# TODO Delete this file after approving /embeddings/document_retrieval_example.py
 
 # [START generativeaionvertexai_sdk_embedding]
 from typing import List, Optional

generative_ai/embedding_batch.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# TODO: Delete this file after approving /embeddings/batch_example.py
 import os
 
 from google.cloud.aiplatform import BatchPredictionJob

generative_ai/embedding_batch_test.py

@@ -11,7 +11,7 @@
 # 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.
-
+# TODO: Delete this file after approving /embeddings/batch_example_test.py
 import os
 
 import backoff

generative_ai/embedding_model_tuning.py

@@ -11,7 +11,7 @@
 # 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.
-
+# TODO Delete this file after approving /embeddings/model_tuning_example.py
 # [START generativeaionvertexai_sdk_embedding]
 import re
 

generative_ai/embedding_model_tuning_test.py

@@ -11,7 +11,7 @@
 # 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.
-
+# TODO Delete this file after approving /embeddings/model_tuning_test.py
 import os
 
 import backoff

generative_ai/embedding_preview.py

@@ -11,7 +11,7 @@
 # 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.
-
+# TODO: Delete this file after approving /embeddings/code_retrieval_example.py
 from __future__ import annotations
 
 # [START generativeaionvertexai_sdk_embedding]

generative_ai/embedding_test.py

@@ -11,7 +11,7 @@
 # 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.
-
+# TODO: Delete this file after approving /embeddings/retrievals_test.py
 import backoff
 
 from google.api_core.exceptions import ResourceExhausted

generative_ai/embeddings/batch_example.py

@@ -0,0 +1,62 @@
+# 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
+
+from google.cloud.aiplatform import BatchPredictionJob
+
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+OUTPUT_URI = os.getenv("GCS_OUTPUT_URI")
+
+
+def embed_text_batch() -> BatchPredictionJob:
+    """Example of how to generate embeddings from text using batch processing.
+
+    Read more: https://cloud.google.com/vertex-ai/generative-ai/docs/embeddings/batch-prediction-genai-embeddings
+    """
+    # [START generativeaionvertexai_embedding_batch]
+    import vertexai
+    from vertexai.preview import language_models
+
+    # TODO(developer): Uncomment and set your project ID
+    # PROJECT_ID = "your-project-id"
+    vertexai.init(project=PROJECT_ID, location="us-central1")
+    input_uri = (
+        "gs://cloud-samples-data/generative-ai/embeddings/embeddings_input.jsonl"
+    )
+    # Format: `"gs://your-bucket-unique-name/directory/` or `bq://project_name.llm_dataset`
+    output_uri = OUTPUT_URI
+
+    textembedding_model = language_models.TextEmbeddingModel.from_pretrained(
+        "textembedding-gecko"
+    )
+
+    batch_prediction_job = textembedding_model.batch_predict(
+        dataset=[input_uri],
+        destination_uri_prefix=output_uri,
+    )
+    print(batch_prediction_job.display_name)
+    print(batch_prediction_job.resource_name)
+    print(batch_prediction_job.state)
+    # Example response:
+    # BatchPredictionJob 2024-09-10 15:47:51.336391
+    # projects/1234567890/locations/us-central1/batchPredictionJobs/123456789012345
+    # JobState.JOB_STATE_SUCCEEDED
+
+    # [END generativeaionvertexai_embedding_batch]
+
+    return batch_prediction_job
+
+
+if __name__ == "__main__":
+    embed_text_batch()

generative_ai/embeddings/batch_example_test.py

@@ -0,0 +1,30 @@
+# 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
+
+import batch_example
+
+from google.api_core.exceptions import ResourceExhausted
+
+import pytest
+
+
+@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
+@pytest.fixture(scope="session", autouse=True)
+def test_embed_text_batch() -> None:
+    os.environ["GCS_OUTPUT_URI"] = "gs://python-docs-samples-tests/"
+    batch_prediction_job = batch_example.embed_text_batch()
+    assert batch_prediction_job

generative_ai/embeddings/code_retrieval_example.py

@@ -0,0 +1,61 @@
+# 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.
+
+from __future__ import annotations
+
+# [START generativeaionvertexai_embedding_code_retrieval]
+from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
+
+MODEL_NAME = "text-embedding-preview-0815"
+DIMENSIONALITY = 256
+
+
+def embed_text(
+    texts: list[str] = ["Retrieve a function that adds two numbers"],
+    task: str = "CODE_RETRIEVAL_QUERY",
+    model_name: str = "text-embedding-preview-0815",
+    dimensionality: int | None = 256,
+) -> list[list[float]]:
+    """Embeds texts with a pre-trained, foundational model."""
+    model = TextEmbeddingModel.from_pretrained(model_name)
+    inputs = [TextEmbeddingInput(text, task) for text in texts]
+    kwargs = dict(output_dimensionality=dimensionality) if dimensionality else {}
+    embeddings = model.get_embeddings(inputs, **kwargs)
+    # Example response:
+    # [[0.025890009477734566, -0.05553026497364044, 0.006374752148985863,...],
+    return [embedding.values for embedding in embeddings]
+
+
+if __name__ == "__main__":
+    # Embeds code block with a pre-trained, foundational model.
+    # Using this function to calculate the embedding for corpus.
+    texts = ["Retrieve a function that adds two numbers"]
+    task = "CODE_RETRIEVAL_QUERY"
+    code_block_embeddings = embed_text(
+        texts=texts, task=task, model_name=MODEL_NAME, dimensionality=DIMENSIONALITY
+    )
+
+    # Embeds code retrieval with a pre-trained, foundational model.
+    # Using this function to calculate the embedding for query.
+    texts = [
+        "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",
+    ]
+    task = "RETRIEVAL_DOCUMENT"
+    code_query_embeddings = embed_text(
+        texts=texts, task=task, model_name=MODEL_NAME, dimensionality=DIMENSIONALITY
+    )
+
+# [END generativeaionvertexai_embedding_code_retrieval]