Skip to content

Commit 353dc4f

Browse files
shilpakancharlashilpakancharla
authored
Add other functions to count_tokens (google-gemini#482)
* Add other functions to count_tokens * Tested count_tokens
1 parent 99e5a11 commit 353dc4f

File tree

1 file changed

+142
-1
lines changed

1 file changed

+142
-1
lines changed

‎samples/rest/count_tokens.sh‎

Lines changed: 142 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,21 @@
11
set -eu
22

3+
SCRIPT_DIR=$(dirname "$0")
4+
MEDIA_DIR=$(realpath ${SCRIPT_DIR}/../../third_party)
5+
6+
TEXT_PATH=${MEDIA_DIR}/poem.txt
7+
IMG_PATH=${MEDIA_DIR}/organ.jpg
8+
AUDIO_PATH=${MEDIA_DIR}/sample.mp3
9+
VIDEO_PATH=${MEDIA_DIR}/Big_Buck_Bunny.mp4
10+
11+
BASE_URL="https://generativelanguage.googleapis.com"
12+
13+
if [[ "$(base64 --version 2>&1)" = *"FreeBSD"* ]]; then
14+
B64FLAGS="--input"
15+
else
16+
B64FLAGS="-w0"
17+
fi
18+
319
echo "[START tokens_text_only]"
420
# [START tokens_text_only]
521
curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY \
@@ -29,4 +45,129 @@ curl https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:co
2945
},
3046
],
3147
}'
32-
# [END tokens_chat]
48+
# [END tokens_chat]
49+
50+
echo "[START tokens_multimodal_image_inline]"
51+
# [START tokens_multimodal_image_inline]
52+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY" \
53+
-H 'Content-Type: application/json' \
54+
-X POST \
55+
-d '{
56+
"contents": [{
57+
"parts":[
58+
{"text": "Tell me about this instrument"},
59+
{
60+
"inline_data": {
61+
"mime_type":"image/jpeg",
62+
"data": "'$(base64 $B64FLAGS $IMG_PATH)'"
63+
}
64+
}
65+
]
66+
}]
67+
}' 2> /dev/null
68+
# [END tokens_multimodal_image_inline]
69+
70+
echo "[START tokens_multimodal_image_file_api]"
71+
# [START tokens_multimodal_image_file_api]
72+
MIME_TYPE=$(file -b --mime-type "${IMG_PATH}")
73+
NUM_BYTES=$(wc -c < "${IMG_PATH}")
74+
DISPLAY_NAME=TEXT
75+
76+
tmp_header_file=upload-header.tmp
77+
78+
# Initial resumable request defining metadata.
79+
# The upload url is in the response headers dump them to a file.
80+
curl "${BASE_URL}/upload/v1beta/files?key=${GOOGLE_API_KEY}" \
81+
-D upload-header.tmp \
82+
-H "X-Goog-Upload-Protocol: resumable" \
83+
-H "X-Goog-Upload-Command: start" \
84+
-H "X-Goog-Upload-Header-Content-Length: ${NUM_BYTES}" \
85+
-H "X-Goog-Upload-Header-Content-Type: ${MIME_TYPE}" \
86+
-H "Content-Type: application/json" \
87+
-d "{'file': {'display_name': '${DISPLAY_NAME}'}}" 2> /dev/null
88+
89+
upload_url=$(grep -i "x-goog-upload-url: " "${tmp_header_file}" | cut -d" " -f2 | tr -d "\r")
90+
rm "${tmp_header_file}"
91+
92+
# Upload the actual bytes.
93+
curl "${upload_url}" \
94+
-H "Content-Length: ${NUM_BYTES}" \
95+
-H "X-Goog-Upload-Offset: 0" \
96+
-H "X-Goog-Upload-Command: upload, finalize" \
97+
--data-binary "@${IMG_PATH}" 2> /dev/null > file_info.json
98+
99+
file_uri=$(jq ".file.uri" file_info.json)
100+
echo file_uri=$file_uri
101+
102+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY" \
103+
-H 'Content-Type: application/json' \
104+
-X POST \
105+
-d '{
106+
"contents": [{
107+
"parts":[
108+
{"text": "Can you tell me about the instruments in this photo?"},
109+
{"file_data":
110+
{"mime_type": "image/jpeg",
111+
"file_uri": '$file_uri'}
112+
}]
113+
}]
114+
}'
115+
# [END tokens_multimodal_image_file_api]
116+
117+
echo "# [START tokens_multimodal_video_audio_file_api]"
118+
# [START tokens_multimodal_video_audio_file_api]
119+
120+
MIME_TYPE=$(file -b --mime-type "${VIDEO_PATH}")
121+
NUM_BYTES=$(wc -c < "${VIDEO_PATH}")
122+
DISPLAY_NAME=VIDEO_PATH
123+
124+
# Initial resumable request defining metadata.
125+
# The upload url is in the response headers dump them to a file.
126+
curl "${BASE_URL}/upload/v1beta/files?key=${GOOGLE_API_KEY}" \
127+
-D upload-header.tmp \
128+
-H "X-Goog-Upload-Protocol: resumable" \
129+
-H "X-Goog-Upload-Command: start" \
130+
-H "X-Goog-Upload-Header-Content-Length: ${NUM_BYTES}" \
131+
-H "X-Goog-Upload-Header-Content-Type: ${MIME_TYPE}" \
132+
-H "Content-Type: application/json" \
133+
-d "{'file': {'display_name': '${DISPLAY_NAME}'}}" 2> /dev/null
134+
135+
upload_url=$(grep -i "x-goog-upload-url: " "${tmp_header_file}" | cut -d" " -f2 | tr -d "\r")
136+
rm "${tmp_header_file}"
137+
138+
# Upload the actual bytes.
139+
curl "${upload_url}" \
140+
-H "Content-Length: ${NUM_BYTES}" \
141+
-H "X-Goog-Upload-Offset: 0" \
142+
-H "X-Goog-Upload-Command: upload, finalize" \
143+
--data-binary "@${VIDEO_PATH}" 2> /dev/null > file_info.json
144+
145+
file_uri=$(jq ".file.uri" file_info.json)
146+
echo file_uri=$file_uri
147+
148+
state=$(jq ".file.state" file_info.json)
149+
echo state=$state
150+
151+
name=$(jq ".file.name" file_info.json)
152+
echo name=$name
153+
154+
while [[ "($state)" = *"PROCESSING"* ]];
155+
do
156+
echo "Processing video..."
157+
sleep 5
158+
# Get the file of interest to check state
159+
curl https://generativelanguage.googleapis.com/v1beta/files/$name > file_info.json
160+
state=$(jq ".file.state" file_info.json)
161+
done
162+
163+
curl "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:countTokens?key=$GOOGLE_API_KEY" \
164+
-H 'Content-Type: application/json' \
165+
-X POST \
166+
-d '{
167+
"contents": [{
168+
"parts":[
169+
{"text": "Describe this video clip"},
170+
{"file_data":{"mime_type": "video/mp4", "file_uri": '$file_uri'}}]
171+
}]
172+
}'
173+
# [END tokens_multimodal_video_audio_file_api]

0 commit comments

Comments
 (0)