fix(generative-ai): Update count tokens sample

generative_ai/embedding_model_tuning_test.py

@@ -43,7 +43,8 @@ def test_tune_embedding_model() -> None:
         credentials=credentials,
     )
     tuning_job = embedding_model_tuning.tune_embedding_model(
-        aiplatform_init.global_config.api_endpoint)
+        aiplatform_init.global_config.api_endpoint
+    )
     try:
         assert tuning_job._status.name != "PIPELINE_STATE_FAILED"
     finally:

generative_ai/gemini_count_token_example.py

@@ -13,21 +13,74 @@
 # limitations under the License.
 
 
-def count_tokens(project_id: str) -> str:
+from vertexai.generative_models import GenerationResponse
+
+
+def count_tokens(project_id: str) -> GenerationResponse:
     # [START generativeaionvertexai_gemini_token_count]
     import vertexai
-
     from vertexai.generative_models import GenerativeModel
 
-    # TODO(developer): Update and un-comment below line
+    # TODO(developer): Update and un-comment below lines
     # project_id = "PROJECT_ID"
 
     vertexai.init(project=project_id, location="us-central1")
 
     model = GenerativeModel(model_name="gemini-1.0-pro-002")
 
-    response = model.count_tokens("Why is sky blue?")
-    print(response)
+    prompt = "Why is the sky blue?"
+
+    # Prompt tokens count
+    response = model.count_tokens(prompt)
+    usage_metadata = response.usage_metadata
+    print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
+
+    # Send text to Gemini
+    response = model.generate_content(prompt)
+
+    # Response tokens count
+    usage_metadata = response.usage_metadata
+    print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
+    print(f"Candidates Token Count: {usage_metadata.candidates_token_count}")
+    print(f"Total Token Count: {usage_metadata.total_token_count}")
 
     # [END generativeaionvertexai_gemini_token_count]
     return response
+
+
+def count_tokens_multimodal(project_id: str) -> GenerationResponse:
+    # [START generativeaionvertexai_gemini_token_count_multimodal]
+    import vertexai
+    from vertexai.generative_models import GenerativeModel, Part
+
+    # TODO(developer): Update and un-comment below lines
+    # project_id = "PROJECT_ID"
+
+    vertexai.init(project=project_id, location="us-central1")
+
+    model = GenerativeModel(model_name="gemini-1.0-pro-002")
+
+    contents = [
+        Part.from_uri(
+            "gs://cloud-samples-data/generative-ai/video/pixel8.mp4",
+            mime_type="video/mp4",
+        ),
+        "Provide a description of the video.",
+    ]
+
+    # Prompt tokens count
+    response = model.count_tokens(contents)
+    usage_metadata = response.usage_metadata
+    print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
+
+    # Send text to Gemini
+    response = model.generate_content(contents)
+    usage_metadata = response.usage_metadata
+
+    # Response tokens count
+    print(f"Prompt Token Count: {usage_metadata.prompt_token_count}")
+    print(f"Candidates Token Count: {usage_metadata.candidates_token_count}")
+    print(f"Total Token Count: {usage_metadata.total_token_count}")
+
+    # [END generativeaionvertexai_gemini_token_count_multimodal]
+    return response

generative_ai/test_gemini_examples.py

@@ -82,6 +82,11 @@ def test_gemini_multi_image_example() -> None:
 def test_gemini_count_token_example() -> None:
     response = gemini_count_token_example.count_tokens(PROJECT_ID)
     assert response
+    assert response.usage_metadata
+
+    response = gemini_count_token_example.count_tokens_multimodal(PROJECT_ID)
+    assert response
+    assert response.usage_metadata
 
 
 def test_gemini_safety_config_example() -> None: