genai: samples for counting tokens in video (#181)

google · Jul 18, 2024 · 0fdd360 · 0fdd360
1 parent ae10a3a
commit 0fdd360
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 6 deletions.
diff --git a/genai/example_test.go b/genai/example_test.go
@@ -531,9 +531,7 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
  }
  defer client.DeleteFile(ctx, file.Name)
 
- fd := genai.FileData{
- URI: file.URI,
- }
+ fd := genai.FileData{URI: file.URI}
  // Call `CountTokens` 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.
@@ -557,6 +555,46 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
 
 }
 
+func ExampleGenerativeModel_CountTokens_videoUploadFile() {
+ ctx := context.Background()
+ client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer client.Close()
+
+ model := client.GenerativeModel("gemini-1.5-flash")
+ prompt := "Tell me about this video"
+ file, err := uploadFile(ctx, client, filepath.Join(testDataDir, "earth.mp4"))
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer client.DeleteFile(ctx, file.Name)
+
+ fd := genai.FileData{URI: file.URI}
+ // Call `CountTokens` 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.
+ tokResp, err := model.CountTokens(ctx, genai.Text(prompt), fd)
+ if err != nil {
+ log.Fatal(err)
+ }
+ fmt.Println("total_tokens:", tokResp.TotalTokens)
+ // ( total_tokens: 1481 )
+
+ resp, err := model.GenerateContent(ctx, genai.Text(prompt), fd)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ fmt.Println("prompt_token_count:", resp.UsageMetadata.PromptTokenCount)
+ fmt.Println("candidates_token_count:", resp.UsageMetadata.CandidatesTokenCount)
+ fmt.Println("total_token_count:", resp.UsageMetadata.TotalTokenCount)
+ // ( prompt_token_count: 1481, candidates_token_count: 43, total_token_count: 1524 )
+
+}
+
 func ExampleGenerativeModel_CountTokens_chat() {
  ctx := context.Background()
  client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))

diff --git a/genai/internal/samples/docs-snippets_test.go b/genai/internal/samples/docs-snippets_test.go
@@ -544,9 +544,7 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
  }
  defer client.DeleteFile(ctx, file.Name)
 
- fd := genai.FileData{
- URI: file.URI,
- }
+ fd := genai.FileData{URI: file.URI}
  // Call `CountTokens` 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.
@@ -570,6 +568,48 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
  // [END tokens_multimodal_image_file_api]
 }
 
+func ExampleGenerativeModel_CountTokens_videoUploadFile() {
+ ctx := context.Background()
+ client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer client.Close()
+
+ // [START tokens_multimodal_video_audio_file_api]
+ model := client.GenerativeModel("gemini-1.5-flash")
+ prompt := "Tell me about this video"
+ file, err := uploadFile(ctx, client, filepath.Join(testDataDir, "earth.mp4"))
+ if err != nil {
+ log.Fatal(err)
+ }
+ defer client.DeleteFile(ctx, file.Name)
+
+ fd := genai.FileData{URI: file.URI}
+ // Call `CountTokens` 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.
+ tokResp, err := model.CountTokens(ctx, genai.Text(prompt), fd)
+ if err != nil {
+ log.Fatal(err)
+ }
+ fmt.Println("total_tokens:", tokResp.TotalTokens)
+ // ( total_tokens: 1481 )
+
+ resp, err := model.GenerateContent(ctx, genai.Text(prompt), fd)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ fmt.Println("prompt_token_count:", resp.UsageMetadata.PromptTokenCount)
+ fmt.Println("candidates_token_count:", resp.UsageMetadata.CandidatesTokenCount)
+ fmt.Println("total_token_count:", resp.UsageMetadata.TotalTokenCount)
+ // ( prompt_token_count: 1481, candidates_token_count: 43, total_token_count: 1524 )
+
+ // [END tokens_multimodal_video_audio_file_api]
+}
+
 func ExampleGenerativeModel_CountTokens_chat() {
  ctx := context.Background()
  client, err := genai.NewClient(ctx, option.WithAPIKey(os.Getenv("GEMINI_API_KEY")))