14
14
# limitations under the License.
15
15
from absl .testing import absltest
16
16
17
- import google .generativeai as genai
18
17
import pathlib
19
18
20
19
media = pathlib .Path (__file__ ).parents [1 ] / "third_party"
23
22
class UnitTests (absltest .TestCase ):
24
23
def test_tokens_context_window (self ):
25
24
# [START tokens_context_window]
25
+ import google .generativeai as genai
26
+
26
27
model_info = genai .get_model ("models/gemini-1.5-flash" )
27
28
28
29
# Returns the "context window" for the model,
@@ -34,6 +35,8 @@ def test_tokens_context_window(self):
34
35
35
36
def test_tokens_text_only (self ):
36
37
# [START tokens_text_only]
38
+ import google .generativeai as genai
39
+
37
40
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
38
41
39
42
prompt = "The quick brown fox jumps over the lazy dog."
@@ -54,6 +57,8 @@ def test_tokens_text_only(self):
54
57
55
58
def test_tokens_chat (self ):
56
59
# [START tokens_chat]
60
+ import google .generativeai as genai
61
+
57
62
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
58
63
59
64
chat = model .start_chat (
@@ -86,6 +91,8 @@ def test_tokens_chat(self):
86
91
87
92
def test_tokens_multimodal_image_inline (self ):
88
93
# [START tokens_multimodal_image_inline]
94
+ import google .generativeai as genai
95
+
89
96
import PIL .Image
90
97
91
98
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -112,6 +119,8 @@ def test_tokens_multimodal_image_inline(self):
112
119
113
120
def test_tokens_multimodal_image_file_api (self ):
114
121
# [START tokens_multimodal_image_file_api]
122
+ import google .generativeai as genai
123
+
115
124
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
116
125
117
126
prompt = "Tell me about this image"
@@ -136,6 +145,8 @@ def test_tokens_multimodal_image_file_api(self):
136
145
137
146
def test_tokens_multimodal_video_audio_file_api (self ):
138
147
# [START tokens_multimodal_video_audio_file_api]
148
+ import google .generativeai as genai
149
+
139
150
import time
140
151
141
152
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -169,6 +180,8 @@ def test_tokens_multimodal_video_audio_file_api(self):
169
180
170
181
def test_tokens_multimodal_pdf_file_api (self ):
171
182
# [START tokens_multimodal_pdf_file_api]
183
+ import google .generativeai as genai
184
+
172
185
model = genai .GenerativeModel ("gemini-1.5-flash" )
173
186
sample_pdf = genai .upload_file (media / "test.pdf" )
174
187
token_count = model .count_tokens (["Give me a summary of this document." , sample_pdf ])
@@ -180,6 +193,8 @@ def test_tokens_multimodal_pdf_file_api(self):
180
193
181
194
def test_tokens_cached_content (self ):
182
195
# [START tokens_cached_content]
196
+ import google .generativeai as genai
197
+
183
198
import time
184
199
185
200
model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -220,6 +235,8 @@ def test_tokens_cached_content(self):
220
235
221
236
def test_tokens_system_instruction (self ):
222
237
# [START tokens_system_instruction]
238
+ import google .generativeai as genai
239
+
223
240
model = genai .GenerativeModel (model_name = "gemini-1.5-flash" )
224
241
225
242
prompt = "The quick brown fox jumps over the lazy dog."
@@ -239,6 +256,8 @@ def test_tokens_system_instruction(self):
239
256
240
257
def test_tokens_tools (self ):
241
258
# [START tokens_tools]
259
+ import google .generativeai as genai
260
+
242
261
model = genai .GenerativeModel (model_name = "gemini-1.5-flash" )
243
262
244
263
prompt = "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
0 commit comments