Skip to content

Commit

Permalink
genai: adjust samples for count tokens (#170)
Browse files Browse the repository at this point in the history 
Aligning with the Python samples per
google-gemini/generative-ai-python#459

Main change is in the way output of snippets is marked
  • Loading branch information
@eliben
eliben authored Jul 12, 2024
1 parent 0700958 commit 19eed92
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 47 deletions.
51 changes: 47 additions & 4 deletions genai/example_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,8 +166,12 @@ func ExampleGenerativeModel_CountTokens_contextWindow() {
if err != nil {
log.Fatal(err)
}
fmt.Println("input_token_limit:", info.InputTokenLimit)
fmt.Println("output_token_limit:", info.OutputTokenLimit)

// Returns the "context window" for the model,
// which is the combined input and output token limits.
fmt.Printf("input_token_limit=%v\n", info.InputTokenLimit)
fmt.Printf("output_token_limit=%v\n", info.OutputTokenLimit)
// ( input_token_limit=30720, output_token_limit=2048 )

}

Expand All @@ -188,6 +192,7 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
}

fmt.Println("total_tokens:", tokResp.TotalTokens)
// ( total_tokens: 10 )

resp, err := model.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
Expand All @@ -197,7 +202,7 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
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: 10, candidates_token_count: 38, total_token_count: 48 )
}

func ExampleGenerativeModel_CountTokens_cachedContent() {
Expand Down Expand Up @@ -226,6 +231,7 @@ func ExampleGenerativeModel_CountTokens_cachedContent() {
log.Fatal(err)
}
fmt.Println("total_tokens:", tokResp.TotalTokens)
// ( total_tokens: 5 )

resp, err := modelWithCache.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
Expand All @@ -236,6 +242,7 @@ func ExampleGenerativeModel_CountTokens_cachedContent() {
fmt.Println("candidates_token_count:", resp.UsageMetadata.CandidatesTokenCount)
fmt.Println("cached_content_token_count:", resp.UsageMetadata.CachedContentTokenCount)
fmt.Println("total_token_count:", resp.UsageMetadata.TotalTokenCount)
// ( prompt_token_count: 33007, candidates_token_count: 39, cached_content_token_count: 33002, total_token_count: 33046 )

}

Expand All @@ -253,12 +260,26 @@ func ExampleGenerativeModel_CountTokens_imageInline() {
if err != nil {
log.Fatal(err)
}

// 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), genai.ImageData("jpeg", imageFile))
if err != nil {
log.Fatal(err)
}
fmt.Println("total_tokens:", tokResp.TotalTokens)
// ( total_tokens: 264 )

resp, err := model.GenerateContent(ctx, genai.Text(prompt), genai.ImageData("jpeg", imageFile))
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: 264, candidates_token_count: 100, total_token_count: 364 )

}

Expand Down Expand Up @@ -286,11 +307,26 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
fd := genai.FileData{
URI: uploadedFile.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: 264 )

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: 264, candidates_token_count: 100, total_token_count: 364 )

}

Expand Down Expand Up @@ -326,9 +362,14 @@ func ExampleGenerativeModel_CountTokens_chat() {
log.Fatal(err)
}

// On the response for SendMessage, use `UsageMetadata` 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`).
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: 25, candidates_token_count: 21, total_token_count: 46 )

}

Expand All @@ -349,6 +390,7 @@ func ExampleGenerativeModel_CountTokens_systemInstruction() {
log.Fatal(err)
}
fmt.Println("total_tokens:", respNoInstruction.TotalTokens)
// ( total_tokens: 10 )

// Same prompt, this time with system instruction
model.SystemInstruction = &genai.Content{
Expand All @@ -359,6 +401,7 @@ func ExampleGenerativeModel_CountTokens_systemInstruction() {
log.Fatal(err)
}
fmt.Println("total_tokens:", respWithInstruction.TotalTokens)
// ( total_tokens: 21 )

}

Expand Down
90 changes: 47 additions & 43 deletions genai/internal/snippets/example_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -167,14 +167,13 @@ func ExampleGenerativeModel_CountTokens_contextWindow() {
if err != nil {
log.Fatal(err)
}
fmt.Println("input_token_limit:", info.InputTokenLimit)
fmt.Println("output_token_limit:", info.OutputTokenLimit)
// [END tokens_context_window]

// [START tokens_context_window_return]
// input_token_limit: 30720
// output_token_limit: 2048
// [END tokens_context_window_return]
// Returns the "context window" for the model,
// which is the combined input and output token limits.
fmt.Printf("input_token_limit=%v\n", info.InputTokenLimit)
fmt.Printf("output_token_limit=%v\n", info.OutputTokenLimit)
// ( input_token_limit=30720, output_token_limit=2048 )
// [END tokens_context_window]
}

func ExampleGenerativeModel_CountTokens_textOnly() {
Expand All @@ -195,6 +194,7 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
}

fmt.Println("total_tokens:", tokResp.TotalTokens)
// ( total_tokens: 10 )

resp, err := model.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
Expand All @@ -204,14 +204,7 @@ func ExampleGenerativeModel_CountTokens_textOnly() {
fmt.Println("prompt_token_count:", resp.UsageMetadata.PromptTokenCount)
fmt.Println("candidates_token_count:", resp.UsageMetadata.CandidatesTokenCount)
fmt.Println("total_token_count:", resp.UsageMetadata.TotalTokenCount)
// [END tokens_text_only]

// [START tokens_text_only_return]
// total_tokens: 10
// prompt_token_count: 10
// candidates_token_count: 38
// total_token_count: 48
// [END tokens_text_only_return]
// ( prompt_token_count: 10, candidates_token_count: 38, total_token_count: 48 )
}

func ExampleGenerativeModel_CountTokens_cachedContent() {
Expand Down Expand Up @@ -241,6 +234,7 @@ func ExampleGenerativeModel_CountTokens_cachedContent() {
log.Fatal(err)
}
fmt.Println("total_tokens:", tokResp.TotalTokens)
// ( total_tokens: 5 )

resp, err := modelWithCache.GenerateContent(ctx, genai.Text(prompt))
if err != nil {
Expand All @@ -251,15 +245,8 @@ func ExampleGenerativeModel_CountTokens_cachedContent() {
fmt.Println("candidates_token_count:", resp.UsageMetadata.CandidatesTokenCount)
fmt.Println("cached_content_token_count:", resp.UsageMetadata.CachedContentTokenCount)
fmt.Println("total_token_count:", resp.UsageMetadata.TotalTokenCount)
// ( prompt_token_count: 33007, candidates_token_count: 39, cached_content_token_count: 33002, total_token_count: 33046 )
// [END tokens_cached_content]

// [START tokens_cached_content_return]
// total_tokens: 5
// prompt_token_count: 33007
// candidates_token_count: 39
// cached_content_token_count: 33002
// total_token_count: 33046
// [END tokens_cached_content_return]
}

func ExampleGenerativeModel_CountTokens_imageInline() {
Expand All @@ -277,17 +264,27 @@ func ExampleGenerativeModel_CountTokens_imageInline() {
if err != nil {
log.Fatal(err)
}

// 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), genai.ImageData("jpeg", imageFile))
if err != nil {
log.Fatal(err)
}
fmt.Println("total_tokens:", tokResp.TotalTokens)
// [END tokens_multimodal_image_inline]
// ( total_tokens: 264 )

resp, err := model.GenerateContent(ctx, genai.Text(prompt), genai.ImageData("jpeg", imageFile))
if err != nil {
log.Fatal(err)
}

// [START tokens_multimodal_image_inline_return]
// total_tokens: 264
// [END tokens_multimodal_image_inline_return]
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: 264, candidates_token_count: 100, total_token_count: 364 )
// [END tokens_multimodal_image_inline]
}

func ExampleGenerativeModel_CountTokens_imageUploadFile() {
Expand Down Expand Up @@ -315,16 +312,27 @@ func ExampleGenerativeModel_CountTokens_imageUploadFile() {
fd := genai.FileData{
URI: uploadedFile.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)
// [END tokens_multimodal_image_file_api]
// ( total_tokens: 264 )

resp, err := model.GenerateContent(ctx, genai.Text(prompt), fd)
if err != nil {
log.Fatal(err)
}

// [START tokens_multimodal_image_file_api_return]
// total_tokens: 264
// [END tokens_multimodal_image_file_api_return]
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: 264, candidates_token_count: 100, total_token_count: 364 )
// [END tokens_multimodal_image_file_api]
}

func ExampleGenerativeModel_CountTokens_chat() {
Expand Down Expand Up @@ -360,16 +368,15 @@ func ExampleGenerativeModel_CountTokens_chat() {
log.Fatal(err)
}

// On the response for SendMessage, use `UsageMetadata` 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`).
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: 25, candidates_token_count: 21, total_token_count: 46 )
// [END tokens_chat]

// [START tokens_chat_return]
// prompt_token_count: 21
// candidates_token_count: 1
// total_token_count: 22
// [END tokens_chat_return]
}

func ExampleGenerativeModel_CountTokens_systemInstruction() {
Expand All @@ -390,6 +397,7 @@ func ExampleGenerativeModel_CountTokens_systemInstruction() {
log.Fatal(err)
}
fmt.Println("total_tokens:", respNoInstruction.TotalTokens)
// ( total_tokens: 10 )

// Same prompt, this time with system instruction
model.SystemInstruction = &genai.Content{
Expand All @@ -400,12 +408,8 @@ func ExampleGenerativeModel_CountTokens_systemInstruction() {
log.Fatal(err)
}
fmt.Println("total_tokens:", respWithInstruction.TotalTokens)
// ( total_tokens: 21 )
// [END tokens_system_instruction]

// [START tokens_system_instruction_return]
// totak_tokens: 10
// totak_tokens: 21
// [END tokens_system_instruction_return]
}

// This example shows how to get a JSON response that conforms to a schema.
Expand Down

0 comments on commit 19eed92

Please sign in to comment.