Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit d426ed8

Browse filesBrowse files
holtskinnergcf-owl-bot[bot]
holtskinner
and
gcf-owl-bot[bot]
authored
feat: Add Vertex AI Grounding samples (GoogleCloudPlatform#11246)
* feat: Add Vertex AI Grounding samples * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Fix lint errors & test errors * Fixed lint error * Removed unused import * Fix lint error * Fixed import order * Change Gemini Model Name to `gemini-1.0-pro` * Update copyright to 2024 --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 9855e23 commit d426ed8
Copy full SHA for d426ed8

File tree

Expand file treeCollapse file tree

4 files changed

+171
-0
lines changed
Filter options
Expand file treeCollapse file tree

4 files changed

+171
-0
lines changed

‎generative_ai/gemini_grounding_example.py

Copy file name to clipboard
+53Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START aiplatform_gemini_grounding]
16+
from typing import Optional
17+
18+
import vertexai
19+
from vertexai.preview.generative_models import (
20+
GenerationResponse,
21+
GenerativeModel,
22+
grounding,
23+
Tool,
24+
)
25+
26+
27+
def generate_text_with_grounding(
28+
project_id: str, location: str, data_store_path: Optional[str] = None
29+
) -> GenerationResponse:
30+
# Initialize Vertex AI
31+
vertexai.init(project=project_id, location=location)
32+
33+
# Load the model
34+
model = GenerativeModel(model_name="gemini-1.0-pro")
35+
36+
# Create Tool for grounding
37+
if data_store_path:
38+
# Use Vertex AI Search data store
39+
# Format: projects/{project_id}/locations/{location}/collections/default_collection/dataStores/{data_store_id}
40+
tool = Tool.from_retrieval(
41+
grounding.Retrieval(grounding.VertexAISearch(datastore=data_store_path))
42+
)
43+
else:
44+
# Use Google Search for grounding (Private Preview)
45+
tool = Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval())
46+
47+
prompt = "What are the price, available colors, and storage size options of a Pixel Tablet?"
48+
response = model.generate_content(prompt, tools=[tool])
49+
50+
print(response)
51+
52+
# [END aiplatform_gemini_grounding]
53+
return response

‎generative_ai/grounding.py

Copy file name to clipboard
+68Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START aiplatform_sdk_grounding]
16+
from typing import Optional
17+
18+
import vertexai
19+
from vertexai.language_models import (
20+
GroundingSource,
21+
TextGenerationModel,
22+
TextGenerationResponse,
23+
)
24+
25+
26+
def grounding(
27+
project_id: str,
28+
location: str,
29+
data_store_location: Optional[str],
30+
data_store_id: Optional[str],
31+
) -> TextGenerationResponse:
32+
"""Grounding example with a Large Language Model"""
33+
34+
vertexai.init(project=project_id, location=location)
35+
36+
# TODO developer - override these parameters as needed:
37+
parameters = {
38+
"temperature": 0.7, # Temperature controls the degree of randomness in token selection.
39+
"max_output_tokens": 256, # Token limit determines the maximum amount of text output.
40+
"top_p": 0.8, # Tokens are selected from most probable to least until the sum of their probabilities equals the top_p value.
41+
"top_k": 40, # A top_k of 1 means the selected token is the most probable among all tokens.
42+
}
43+
44+
model = TextGenerationModel.from_pretrained("text-bison@002")
45+
46+
if data_store_id and data_store_location:
47+
# Use Vertex AI Search data store
48+
grounding_source = GroundingSource.VertexAISearch(
49+
data_store_id=data_store_id, location=data_store_location
50+
)
51+
else:
52+
# Use Google Search for grounding (Private Preview)
53+
grounding_source = GroundingSource.WebSearch()
54+
55+
response = model.predict(
56+
"What are the price, available colors, and storage size options of a Pixel Tablet?",
57+
grounding_source=grounding_source,
58+
**parameters,
59+
)
60+
print(f"Response from Model: {response.text}")
61+
print(f"Grounding Metadata: {response.grounding_metadata}")
62+
# [END aiplatform_sdk_grounding]
63+
64+
return response
65+
66+
67+
if __name__ == "__main__":
68+
grounding()

‎generative_ai/grounding_test.py

Copy file name to clipboard
+38Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
# Copyright 2024 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
import backoff
18+
from google.api_core.exceptions import ResourceExhausted
19+
20+
import grounding
21+
22+
23+
_PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
24+
_LOCATION = "us-central1"
25+
26+
27+
@backoff.on_exception(backoff.expo, ResourceExhausted, max_time=10)
28+
def test_grounding() -> None:
29+
data_store_id = "test-search-engine_1689960780551"
30+
response = grounding.grounding(
31+
project_id=_PROJECT_ID,
32+
location=_LOCATION,
33+
data_store_location="global",
34+
data_store_id=data_store_id,
35+
)
36+
assert response
37+
assert response.text
38+
assert response.grounding_metadata

‎generative_ai/test_gemini_examples.py

Copy file name to clipboardExpand all lines: generative_ai/test_gemini_examples.py
+12Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import gemini_chat_example
2020
import gemini_count_token_example
21+
import gemini_grounding_example
2122
import gemini_guide_example
2223
import gemini_multi_image_example
2324
import gemini_pro_basic_example
@@ -129,3 +130,14 @@ def test_gemini_chat_example() -> None:
129130
text = text.lower()
130131
assert len(text) > 0
131132
assert any([_ in text for _ in ("hi", "hello", "greeting")])
133+
134+
135+
def test_gemini_grounding_example() -> None:
136+
# Test Vertex AI Search Grounding
137+
# Unable to test Google Search grounding due to allowlist restrictions.
138+
data_store_id = "test-search-engine_1689960780551"
139+
data_store_path = f"projects/{PROJECT_ID}/locations/{LOCATION}/collections/default_collection/dataStores/{data_store_id}"
140+
response = gemini_grounding_example.generate_text_with_grounding(
141+
PROJECT_ID, LOCATION, data_store_path=data_store_path
142+
)
143+
assert response

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.