1414# limitations under the License.
1515from absl .testing import absltest
1616
17- import google .generativeai as genai
1817import pathlib
1918
2019media = pathlib .Path (__file__ ).parents [1 ] / "third_party"
2322class UnitTests (absltest .TestCase ):
2423 def test_tokens_context_window (self ):
2524 # [START tokens_context_window]
25+ import google .generativeai as genai
26+
2627 model_info = genai .get_model ("models/gemini-1.5-flash" )
2728
2829 # Returns the "context window" for the model,
@@ -34,6 +35,8 @@ def test_tokens_context_window(self):
3435
3536 def test_tokens_text_only (self ):
3637 # [START tokens_text_only]
38+ import google .generativeai as genai
39+
3740 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
3841
3942 prompt = "The quick brown fox jumps over the lazy dog."
@@ -54,6 +57,8 @@ def test_tokens_text_only(self):
5457
5558 def test_tokens_chat (self ):
5659 # [START tokens_chat]
60+ import google .generativeai as genai
61+
5762 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
5863
5964 chat = model .start_chat (
@@ -86,6 +91,8 @@ def test_tokens_chat(self):
8691
8792 def test_tokens_multimodal_image_inline (self ):
8893 # [START tokens_multimodal_image_inline]
94+ import google .generativeai as genai
95+
8996 import PIL .Image
9097
9198 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -112,6 +119,8 @@ def test_tokens_multimodal_image_inline(self):
112119
113120 def test_tokens_multimodal_image_file_api (self ):
114121 # [START tokens_multimodal_image_file_api]
122+ import google .generativeai as genai
123+
115124 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
116125
117126 prompt = "Tell me about this image"
@@ -136,6 +145,8 @@ def test_tokens_multimodal_image_file_api(self):
136145
137146 def test_tokens_multimodal_video_audio_file_api (self ):
138147 # [START tokens_multimodal_video_audio_file_api]
148+ import google .generativeai as genai
149+
139150 import time
140151
141152 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -169,6 +180,8 @@ def test_tokens_multimodal_video_audio_file_api(self):
169180
170181 def test_tokens_multimodal_pdf_file_api (self ):
171182 # [START tokens_multimodal_pdf_file_api]
183+ import google .generativeai as genai
184+
172185 model = genai .GenerativeModel ("gemini-1.5-flash" )
173186 sample_pdf = genai .upload_file (media / "test.pdf" )
174187 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):
180193
181194 def test_tokens_cached_content (self ):
182195 # [START tokens_cached_content]
196+ import google .generativeai as genai
197+
183198 import time
184199
185200 model = genai .GenerativeModel ("models/gemini-1.5-flash" )
@@ -220,6 +235,8 @@ def test_tokens_cached_content(self):
220235
221236 def test_tokens_system_instruction (self ):
222237 # [START tokens_system_instruction]
238+ import google .generativeai as genai
239+
223240 model = genai .GenerativeModel (model_name = "gemini-1.5-flash" )
224241
225242 prompt = "The quick brown fox jumps over the lazy dog."
@@ -239,6 +256,8 @@ def test_tokens_system_instruction(self):
239256
240257 def test_tokens_tools (self ):
241258 # [START tokens_tools]
259+ import google .generativeai as genai
260+
242261 model = genai .GenerativeModel (model_name = "gemini-1.5-flash" )
243262
244263 prompt = "I have 57 cats, each owns 44 mittens, how many mittens is that in total?"
0 commit comments