Skip to content

Commit f53e2a4

Browse files
holtskinnerholtskinner
authored
chore: Update Model Optimizer Notebook to use Gen AI SDK with Model Selection Config (#1982)
1 parent 985fce4 commit f53e2a4

1 file changed

Lines changed: 93 additions & 94 deletions

File tree

‎gemini/model-optimizer/intro_model_optimizer.ipynb‎

Lines changed: 93 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -107,7 +107,7 @@
107107
"id": "MFgnm8KDNJQq"
108108
},
109109
"source": [
110-
"Model Optimizer intelligently routes user queries across models (Pro, Flash, etc.) for resource optimization. This notebook provides an introduction on how to use Model Optimizer. It demonstrates how to send different types of queries (single-turn, multi-turn, and function-call) to the model, and how to configure routing options to control cost and quality.\n",
110+
"Model Optimizer intelligently routes user queries across models (Pro, Flash, etc.) for resource optimization. This notebook provides an introduction on how to use Model Optimizer with the Google Gen AI SDK.\n",
111111
"\n",
112112
"Learn more about [Model Optimizer](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/vertex-ai-model-optimizer)."
113113
]
@@ -125,7 +125,7 @@
125125
"- Send prompt queries to Model Optimizer.\n",
126126
"- Configure routing preferences to optimize quality or cost.\n",
127127
"- Utilize Model Optimizer for single-turn and multi-turn conversations.\n",
128-
"- Execute function calls through Model Optimizer.\n"
128+
"- Execute function calls through Model Optimizer."
129129
]
130130
},
131131
{
@@ -157,7 +157,9 @@
157157
"id": "2qQaXxTIOx5N"
158158
},
159159
"source": [
160-
"### Install libraries"
160+
"### Install Google Gen AI SDK\n",
161+
"\n",
162+
"Learn more about [Google Gen AI SDK](https://cloud.google.com/vertex-ai/generative-ai/docs/sdks/overview)."
161163
]
162164
},
163165
{
@@ -168,35 +170,7 @@
168170
},
169171
"outputs": [],
170172
"source": [
171-
"%pip install --upgrade --quiet \"google-cloud-aiplatform>=1.88.0\""
172-
]
173-
},
174-
{
175-
"cell_type": "markdown",
176-
"metadata": {
177-
"id": "10YwtoGTifpr"
178-
},
179-
"source": [
180-
"### Import libraries"
181-
]
182-
},
183-
{
184-
"cell_type": "code",
185-
"execution_count": null,
186-
"metadata": {
187-
"id": "tSeSuXHNin_K"
188-
},
189-
"outputs": [],
190-
"source": [
191-
"from IPython.display import Markdown, display\n",
192-
"from vertexai.generative_models import (\n",
193-
" Content,\n",
194-
" FunctionDeclaration,\n",
195-
" GenerationConfig,\n",
196-
" GenerativeModel,\n",
197-
" Part,\n",
198-
" Tool,\n",
199-
")"
173+
"%pip install --upgrade --user --quiet google-genai"
200174
]
201175
},
202176
{
@@ -256,52 +230,49 @@
256230
"\n",
257231
"LOCATION = os.environ.get(\"GOOGLE_CLOUD_REGION\", \"us-central1\")\n",
258232
"\n",
259-
"import vertexai\n",
233+
"from google import genai\n",
260234
"\n",
261-
"vertexai.init(project=PROJECT_ID, location=LOCATION)"
235+
"client = genai.Client(vertexai=True, project=PROJECT_ID, location=LOCATION)"
262236
]
263237
},
264238
{
265239
"cell_type": "markdown",
266240
"metadata": {
267-
"id": "F9t-K8prqP2E"
241+
"id": "10YwtoGTifpr"
268242
},
269243
"source": [
270-
"### Choose Routing Option\n",
271-
"\n",
272-
"Model Optimizer intelligently routes user queries based on user-defined cost and quality preferences. This dynamic routing capability enables users to optimize their interactions according to their specific needs."
244+
"### Import libraries"
273245
]
274246
},
275247
{
276248
"cell_type": "code",
277249
"execution_count": null,
278250
"metadata": {
279-
"id": "qeYksAdddeWl"
251+
"id": "tSeSuXHNin_K"
280252
},
281253
"outputs": [],
282254
"source": [
283-
"feature_selection_preference_map = {\n",
284-
" \"PRIORITIZE_COST\": GenerationConfig.ModelConfig.FeatureSelectionPreference.PRIORITIZE_COST,\n",
285-
" \"BALANCED\": GenerationConfig.ModelConfig.FeatureSelectionPreference.BALANCED,\n",
286-
" \"PRIORITIZE_QUALITY\": GenerationConfig.ModelConfig.FeatureSelectionPreference.PRIORITIZE_QUALITY,\n",
287-
"}\n",
288-
"\n",
289-
"preference = \"BALANCED\" # @param ['PRIORITIZE_COST', 'BALANCED', 'PRIORITIZE_QUALITY']\n",
290-
"\n",
291-
"feature_selection = feature_selection_preference_map[preference]\n",
292-
"\n",
293-
"model_config = GenerationConfig.ModelConfig(\n",
294-
" feature_selection_preference=feature_selection,\n",
255+
"from IPython.display import Markdown, display\n",
256+
"from google.genai.types import (\n",
257+
" FeatureSelectionPreference,\n",
258+
" FunctionDeclaration,\n",
259+
" GenerateContentConfig,\n",
260+
" ModelSelectionConfig,\n",
261+
" Part,\n",
262+
" Tool,\n",
263+
" UserContent,\n",
295264
")"
296265
]
297266
},
298267
{
299268
"cell_type": "markdown",
300269
"metadata": {
301-
"id": "2vflN6QxxE_D"
270+
"id": "F9t-K8prqP2E"
302271
},
303272
"source": [
304-
"### Load Model Optimizer Model & set System Instruction"
273+
"### Load Model & set System Instruction\n",
274+
"\n",
275+
"`model-optimizer-exp-04-09` is the model used for Model Optimizer."
305276
]
306277
},
307278
{
@@ -312,10 +283,31 @@
312283
},
313284
"outputs": [],
314285
"source": [
315-
"system_instruction = \"\" # @param {'type': 'string'}\n",
286+
"MODEL_ID = \"model-optimizer-exp-04-09\"\n",
287+
"system_instruction = \"\" # @param {'type': 'string'}"
288+
]
289+
},
290+
{
291+
"cell_type": "markdown",
292+
"metadata": {
293+
"id": "9669b4ec9703"
294+
},
295+
"source": [
296+
"### Choose Routing Option\n",
316297
"\n",
317-
"model = GenerativeModel(\n",
318-
" \"model-optimizer-exp-04-09\", system_instruction=[system_instruction]\n",
298+
"Model Optimizer intelligently routes user queries based on user-defined cost and quality preferences. This dynamic routing capability enables users to optimize their interactions according to their specific needs."
299+
]
300+
},
301+
{
302+
"cell_type": "code",
303+
"execution_count": null,
304+
"metadata": {
305+
"id": "653f0e04f4b2"
306+
},
307+
"outputs": [],
308+
"source": [
309+
"model_selection_config = ModelSelectionConfig(\n",
310+
" feature_selection_preference=FeatureSelectionPreference.BALANCED # Options: PRIORITIZE_QUALITY, BALANCED, PRIORITIZE_COST\n",
319311
")"
320312
]
321313
},
@@ -338,18 +330,25 @@
338330
},
339331
"outputs": [],
340332
"source": [
341-
"def generate_content(prompt, stream: bool = False):\n",
333+
"def generate_content(prompt, stream: bool = False) -> None:\n",
342334
" if stream:\n",
343-
" for chunk in model.generate_content(\n",
335+
" for chunk in client.models.generate_content_stream(\n",
336+
" model=MODEL_ID,\n",
344337
" contents=prompt,\n",
345-
" generation_config=GenerationConfig(model_config=model_config),\n",
346-
" stream=True,\n",
338+
" config=GenerateContentConfig(\n",
339+
" system_instruction=system_instruction,\n",
340+
" model_selection_config=model_selection_config,\n",
341+
" ),\n",
347342
" ):\n",
348343
" print(chunk.text, end=\"\")\n",
349344
" else:\n",
350-
" response = model.generate_content(\n",
345+
" response = client.models.generate_content(\n",
346+
" model=MODEL_ID,\n",
351347
" contents=prompt,\n",
352-
" generation_config=GenerationConfig(model_config=model_config),\n",
348+
" config=GenerateContentConfig(\n",
349+
" system_instruction=system_instruction,\n",
350+
" model_selection_config=model_selection_config,\n",
351+
" ),\n",
353352
" )\n",
354353
" display(Markdown(response.text))"
355354
]
@@ -374,7 +373,7 @@
374373
"outputs": [],
375374
"source": [
376375
"prompt = \"Write a poem\" # @param {'type': 'string'}\n",
377-
"generate_content(prompt, stream=False)"
376+
"generate_content(prompt)"
378377
]
379378
},
380379
{
@@ -390,7 +389,7 @@
390389
"There are two ways to structure multi-turn queries:\n",
391390
"\n",
392391
"- **List of text prompts:** A simple list of strings, where each string represents a turn in the conversation. This format is suitable for basic multi-turn scenarios.\n",
393-
"- **List of Content objects:** A more structured format using Content objects, where you can explicitly define the role (user or model) and parts of each turn. This format provides greater control and clarity, especially when dealing with complex conversations."
392+
"- **List of `Content` objects:** A more structured format using `Content` objects, where you can explicitly define the role (user or model) and parts of each turn. This format provides greater control and clarity, especially when dealing with complex conversations."
394393
]
395394
},
396395
{
@@ -401,11 +400,11 @@
401400
},
402401
"outputs": [],
403402
"source": [
404-
"prompts = [\n",
403+
"prompt = [\n",
405404
" \"What is x multiplied by 2?\",\n",
406405
" \"x = 42\",\n",
407406
"]\n",
408-
"generate_content(prompt, stream=False)"
407+
"generate_content(prompt)"
409408
]
410409
},
411410
{
@@ -425,12 +424,12 @@
425424
},
426425
"outputs": [],
427426
"source": [
428-
"prompts = [\n",
429-
" Content(role=\"user\", parts=[Part.from_text(text=\"Hello\")]),\n",
430-
" Content(role=\"user\", parts=[Part.from_text(text=\"How are you?\")]),\n",
431-
" Content(role=\"user\", parts=[Part.from_text(text=\"Why is the sky blue?\")]),\n",
427+
"prompt = [\n",
428+
" UserContent(parts=[Part.from_text(text=\"Hello\")]),\n",
429+
" UserContent(parts=[Part.from_text(text=\"How are you?\")]),\n",
430+
" UserContent(parts=[Part.from_text(text=\"Why is the sky blue?\")]),\n",
432431
"]\n",
433-
"generate_content(prompt, stream=False)"
432+
"generate_content(prompt)"
434433
]
435434
},
436435
{
@@ -462,11 +461,11 @@
462461
" Returns:\n",
463462
" The weather information as a dict.\n",
464463
" \"\"\"\n",
465-
" return dict(\n",
466-
" location=location,\n",
467-
" unit=unit,\n",
468-
" weather=\"Super nice, but maybe a bit hot.\",\n",
469-
" )\n",
464+
" return {\n",
465+
" \"location\": location,\n",
466+
" \"unit\": unit,\n",
467+
" \"weather\": \"Super nice, but maybe a bit hot.\",\n",
468+
" }\n",
470469
"\n",
471470
"\n",
472471
"_REQUEST_FUNCTION_PARAMETER_SCHEMA_STRUCT = {\n",
@@ -514,53 +513,53 @@
514513
" response=_REQUEST_FUNCTION_RESPONSE_SCHEMA_STRUCT,\n",
515514
")\n",
516515
"\n",
516+
"\n",
517517
"prompt = \"What is the weather like in Boston?\"\n",
518518
"\n",
519519
"weather_tool = Tool(function_declarations=[get_current_weather_func])\n",
520520
"\n",
521-
"\n",
522-
"response = model.generate_content(\n",
523-
" contents=prompt, generation_config=GenerationConfig(model_config=model_config)\n",
524-
")\n",
525-
"\n",
526-
"response = model.generate_content(\n",
521+
"response = client.models.generate_content(\n",
522+
" model=MODEL_ID,\n",
527523
" contents=prompt,\n",
528-
" generation_config=GenerationConfig(model_config=model_config),\n",
529-
" tools=[weather_tool],\n",
524+
" config=GenerateContentConfig(\n",
525+
" tools=[weather_tool],\n",
526+
" model_selection_config=model_selection_config,\n",
527+
" system_instruction=system_instruction,\n",
528+
" ),\n",
530529
")\n",
531530
"\n",
532-
"print(\"Called function: \", response.candidates[0].content.parts[0].function_call.name)"
531+
"print(\"Called function: \", response.function_calls)"
533532
]
534533
},
535534
{
536535
"cell_type": "markdown",
537536
"metadata": {
538-
"id": "iUmPbSxeBg2T"
537+
"id": "234295d42fed"
539538
},
540539
"source": [
541-
"To preview more details:\n"
540+
"To preview more details:"
542541
]
543542
},
544543
{
545544
"cell_type": "code",
546545
"execution_count": null,
547546
"metadata": {
548-
"id": "okMzJBCeBsBM"
547+
"id": "592cff825422"
549548
},
550549
"outputs": [],
551550
"source": [
552551
"function_map = {\"get_current_weather\": get_current_weather}\n",
553552
"\n",
554553
"function_response_parts = []\n",
555-
"for part in response.candidates[0].content.parts:\n",
556-
" function_call = part.function_call\n",
554+
"for function_call in response.function_calls:\n",
557555
" function = function_map[function_call.name]\n",
558556
" function_result = function(**function_call.args)\n",
559-
" function_response_part = Part.from_function_response(\n",
560-
" name=function_call.name,\n",
561-
" response=function_result,\n",
557+
" function_response_parts.append(\n",
558+
" Part.from_function_response(\n",
559+
" name=function_call.name,\n",
560+
" response=function_result,\n",
561+
" )\n",
562562
" )\n",
563-
" function_response_parts.append(function_response_part)\n",
564563
"\n",
565564
"print(function_response_parts)"
566565
]

0 commit comments

Comments
 (0)