1
+ # -*- coding: utf-8 -*-
2
+ # Copyright 2023 Google LLC
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ from absl .testing import absltest
16
+
17
+ import google
18
+ import google .generativeai as genai
19
+ import pathlib
20
+
21
+ class UnitTests (absltest .TestCase ):
22
+ def test_embed_content (self ):
23
+ # [START embed_content]
24
+ text = "Hello World!"
25
+ result = genai .embed_content (model = "models/text-embedding-004" ,
26
+ content = text )
27
+ print (result ['embedding' ])
28
+ print ()
29
+ # [END embed_content]
30
+
31
+ def batch_embed_content (self ):
32
+ # [START batch_embed_content]
33
+ texts = ["What is the meaning of life?" ,
34
+ "How much wood would a woodchuck chuck?" ,
35
+ "How does the brain work?" ]
36
+ result = genai .embed_content (model = "models/text-embedding-004" ,
37
+ content = texts )
38
+ print (result )
39
+ print ()
40
+ # [END batch_embed_content]
41
+
42
+ if __name__ == "__main__" :
43
+ absltest .main ()
0 commit comments