Update count_tokens.py (#459)

- integrated returns into main snippet
- updated code comments
- pulled text of prompts out of the requests to generate_content

Author: rachelsaunders

samples/count_tokens.py

@@ -24,16 +24,14 @@ class UnitTests(absltest.TestCase):
     def test_tokens_context_window(self):
         # [START tokens_context_window]
         model_info = genai.get_model("models/gemini-1.0-pro-001")
-        # Returns the "context window" for the model (the combined input and output token limits)
+
+        # Returns the "context window" for the model,
+        # which is the combined input and output token limits.
         print(f"{model_info.input_token_limit=}")
         print(f"{model_info.output_token_limit=}")
+        # ( input_token_limit=30720, output_token_limit=2048 )
         # [END tokens_context_window]
 
-        # [START tokens_context_window_return]
-        # input_token_limit=30720
-        # output_token_limit=2048
-        # [END tokens_context_window_return]
-
     def test_tokens_text_only(self):
         # [START tokens_text_only]
         model = genai.GenerativeModel("models/gemini-1.5-flash")
@@ -42,22 +40,18 @@ def test_tokens_text_only(self):
 
         # Call `count_tokens` to get the input token count (`total_tokens`).
         print("total_tokens: ", model.count_tokens(prompt))
+        # ( total_tokens: 10 )
 
         response = model.generate_content(prompt)
 
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+        # On the response for `generate_content`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the combined token count (`total_token_count`).
         print(response.usage_metadata)
+        # ( prompt_token_count: 11, candidates_token_count: 73, total_token_count: 84 )
         # [END tokens_text_only]
 
-        # [START tokens_text_only_return]
-        # total_tokens: total_tokens: 10
-        #
-        # prompt_token_count: 11
-        # candidates_token_count: 73
-        # total_token_count: 84
-        # [END tokens_text_only_return]
-
     def test_tokens_chat(self):
         # [START tokens_chat]
         model = genai.GenerativeModel("models/gemini-1.5-flash")
@@ -70,30 +64,26 @@ def test_tokens_chat(self):
         )
         # Call `count_tokens` to get the input token count (`total_tokens`).
         print(model.count_tokens(chat.history))
+        # ( total_tokens: 10 )
 
         response = chat.send_message(
             "In one sentence, explain how a computer works to a young child."
         )
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+    
+        # On the response for `send_message`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the combined token count (`total_token_count`).
         print(response.usage_metadata)
+        # ( prompt_token_count: 25, candidates_token_count: 21, total_token_count: 46 )
 
-        # TODO add comment...
         from google.generativeai.types.content_types import to_contents
 
+        # You can call `count_tokens` on the combined history and content of the next turn.
         print(model.count_tokens(chat.history + to_contents("What is the meaning of life?")))
+        # ( total_tokens: 56 )
         # [END tokens_chat]
 
-        # [START tokens_chat_return]
-        # total_tokens: 10
-        #
-        # prompt_token_count: 25
-        # candidates_token_count: 21
-        # total_token_count: 46
-        #
-        # total_tokens: 56
-        # [END tokens_chat_return]
-
     def test_tokens_multimodal_image_inline(self):
         # [START tokens_multimodal_image_inline]
         import PIL.Image
@@ -103,52 +93,47 @@ def test_tokens_multimodal_image_inline(self):
         prompt = "Tell me about this image"
         your_image_file = PIL.Image.open("image.jpg")
 
-        # Call `count_tokens` to get input token count of the combined text and file (`total_tokens`).
-        # An image's display size does not affect its token count.
-        # Optionally, you can call `count_tokens` for the prompt and file separately.
+        # Call `count_tokens` to get the input token count
+        # of the combined text and file (`total_tokens`).
+        # An image's display or file size does not affect its token count.
+        # Optionally, you can call `count_tokens` for the text and file separately.
         print(model.count_tokens([prompt, your_image_file]))
+        # ( total_tokens: 263 )
 
         response = model.generate_content([prompt, your_image_file])
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+
+        # On the response for `generate_content`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the combined token count (`total_token_count`).
         print(response.usage_metadata)
+        # ( prompt_token_count: 264, candidates_token_count: 80, total_token_count: 345 )
         # [END tokens_multimodal_image_inline]
 
-        # [START tokens_multimodal_image_inline_return]
-        # total_tokens: 263
-        #
-        # prompt_token_count: 264
-        # candidates_token_count: 81
-        # total_token_count: 345
-        # [END tokens_multimodal_image_inline_return]
-
     def test_tokens_multimodal_image_file_api(self):
         # [START tokens_multimodal_image_file_api]
         model = genai.GenerativeModel("models/gemini-1.5-flash")
 
         prompt = "Tell me about this image"
         your_image_file = genai.upload_file(path="image.jpg")
 
-        # Call `count_tokens` to get input token count of the combined text and file (`total_tokens`).
-        # An image's display size does not affect its token count.
-        # Optionally, you can call `count_tokens` for the prompt and file separately.
+        # Call `count_tokens` to get the input token count
+        # of the combined text and file (`total_tokens`).
+        # An image's display or file size does not affect its token count.
+        # Optionally, you can call `count_tokens` for the text and file separately.
         print(model.count_tokens([prompt, your_image_file]))
+        # ( total_tokens: 263 )
 
         response = model.generate_content([prompt, your_image_file])
         response.text
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+        # On the response for `generate_content`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the combined token count (`total_token_count`).
         print(response.usage_metadata)
+        # ( prompt_token_count: 264, candidates_token_count: 80, total_token_count: 345 )
         # [END tokens_multimodal_image_file_api]
 
-        # [START tokens_multimodal_image_file_api_return]
-        # total_tokens: 263
-        #
-        # prompt_token_count: 264
-        # candidates_token_count: 80
-        # total_token_count: 344
-        # [END tokens_multimodal_image_file_api_return]
-
     def test_tokens_multimodal_video_audio_file_api(self):
         # [START tokens_multimodal_video_audio_file_api]
         import time
@@ -164,28 +149,24 @@ def test_tokens_multimodal_video_audio_file_api(self):
             time.sleep(5)
             your_file = genai.get_file(your_file.name)
 
-        # Call `count_tokens` to get input token count of the combined text and file (`total_tokens`).
+        # Call `count_tokens` to get the input token count
+        # of the combined text and video/audio file (`total_tokens`).
         # A video or audio file is converted to tokens at a fixed rate of tokens per second.
-        # Optionally, you can call `count_tokens` for the prompt and file separately.
+        # Optionally, you can call `count_tokens` for the text and file separately.
         print(model.count_tokens([prompt, your_file]))
+        # ( total_tokens: 300 )
 
         response = model.generate_content([prompt, your_file])
 
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+        # On the response for `generate_content`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the combined token count (`total_token_count`).
         print(response.usage_metadata)
+        # ( prompt_token_count: 301, candidates_token_count: 60, total_token_count: 361 )
 
         # [END tokens_multimodal_video_audio_file_api]
 
-        # [START tokens_multimodal_video_audio_file_api_return]
-        # processing video...
-        # total_tokens: 300
-        #
-        # prompt_token_count: 301
-        # candidates_token_count: 60
-        # total_token_count: 361
-        # [END tokens_multimodal_video_audio_file_api_return]
-
     def test_tokens_cached_content(self):
         # [START tokens_cached_content]
         import time
@@ -196,71 +177,63 @@ def test_tokens_cached_content(self):
 
         cache = genai.caching.CachedContent.create(
             model="models/gemini-1.5-flash-001",
-            # You could set the system_instruction and tools
+            # You can set the system_instruction and tools
             system_instruction=None,
             tools=None,
             contents=["Here the Apollo 11 transcript:", your_file],
         )
 
         model = genai.GenerativeModel.from_cached_content(cache)
 
-        # Call `count_tokens` to get input token count of the combined text and file (`total_tokens`).
-        # A video or audio file is converted to tokens at a fixed rate of tokens per second.
-        # Optionally, you can call `count_tokens` for the prompt and file separately.
         prompt = "Please give a short summary of this file."
+
+        # Call `count_tokens` to get input token count
+        # of the combined text and file (`total_tokens`).
+        # A video or audio file is converted to tokens at a fixed rate of tokens per second.
+        # Optionally, you can call `count_tokens` for the text and file separately.
         print(model.count_tokens(prompt))
+        # ( total_tokens: 9 )
 
         response = model.generate_content(prompt)
-        # Use `usage_metadata` to get both input and output token counts
-        # (`prompt_token_count` and `candidates_token_count`, respectively).
+    
+        # On the response for `generate_content`, use `usage_metadata`
+        # to get separate input and output token counts
+        # (`prompt_token_count` and `candidates_token_count`, respectively),
+        # as well as the cached content token count and the combined total token count.
         print(response.usage_metadata)
+        # ( prompt_token_count: 323393, cached_content_token_count: 323383, candidates_token_count: 64)
+        # ( total_token_count: 323457 )
 
         cache.delete()
         # [END tokens_cached_content]
 
-        # [START tokens_cached_content_return]
-        # total_tokens: 9
-        #
-        # prompt_token_count: 323393
-        # cached_content_token_count: 323383
-        # candidates_token_count: 64
-        # total_token_count: 323457
-        # [END tokens_cached_content_return]
-
     def test_tokens_system_instruction(self):
         # [START tokens_system_instruction]
         model = genai.GenerativeModel(model_name="gemini-1.5-flash")
 
-        # The total token count includes everything sent to the generate_content request.
-        print(model.count_tokens("The quick brown fox jumps over the lazy dog."))
+        prompt="The quick brown fox jumps over the lazy dog."
+
+        print(model.count_tokens(prompt))
         # total_tokens: 10
 
         model = genai.GenerativeModel(
             model_name="gemini-1.5-flash", system_instruction="You are a cat. Your name is Neko."
         )
 
-        # The total token count includes everything sent to the generate_content request.
+        # The total token count includes everything sent to the `generate_content` request.
         # When you use system instructions, the total token count increases.
-        print(model.count_tokens("The quick brown fox jumps over the lazy dog."))
+        print(model.count_tokens(prompt))
+        # ( total_tokens: 21 )
         # [END tokens_system_instruction]
 
-        # [START tokens_system_instruction_return]
-        # total_tokens: 10
-        #
-        # total_tokens: 21
-        # [END tokens_system_instruction_return]
-
     def test_tokens_tools(self):
         # [START tokens_tools]
         model = genai.GenerativeModel(model_name="gemini-1.5-flash")
 
-        # The total token count includes everything sent to the generate_content request.
-        print(
-            model.count_tokens(
-                "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
-            )
-        )
-        # total_tokens: 10
+        prompt="I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
+
+        print(model.count_tokens(prompt))
+        # ( total_tokens: 22 )
 
         def add(a: float, b: float):
             """returns a + b."""
@@ -282,19 +255,12 @@ def divide(a: float, b: float):
             "models/gemini-1.5-flash-001", tools=[add, subtract, multiply, divide]
         )
 
-        print(
-            model.count_tokens(
-                "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
-            )
-        )
+        # The total token count includes everything sent to the `generate_content` request.
+        # When you use tools (like function calling), the total token count increases.
+        print(model.count_tokens(prompt))
+        # ( total_tokens: 206 )
         # [END tokens_tools]
 
-        # [START tokens_tools_return]
-        # total_tokens: 22
-        #
-        # total_tokens: 206
-        # [END tokens_tools_return]
-
 
 if __name__ == "__main__":
     absltest.main()