-
Notifications
You must be signed in to change notification settings - Fork 439
Description
Description of the feature request:
Can i use this code to export the FunctionGemma model?
from ai_edge_torch.generative.examples.gemma3 import gemma3
from ai_edge_torch.generative.utilities import converter
from ai_edge_torch.generative.utilities.export_config import ExportConfig
from ai_edge_torch.generative.layers import kv_cache
pytorch_model = gemma3.build_model_270m("PATH_TO_HF_MODEL
")
If you are using Gemma 3 1B
#pytorch_model = gemma3.build_model_1b("PATH_TO_HF_MODEL
")
export_config = ExportConfig()
export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED
export_config.mask_as_input = True
converter.convert_to_tflite(
pytorch_model,
output_path="OUTPUT_DIR_PATH
",
output_name_prefix="my-gemma3",
prefill_seq_len=2048,
kv_cache_max_len=4096,
quantize="dynamic_int8",
export_config=export_config,
)
this code from https://ai.google.dev/gemma/docs/conversions/hf-to-mediapipe-task
if that's true,how can i test the model like this
weather_function_schema = {
"type": "function",
"function": {
"name": "get_current_temperature",
"description": "Gets the current temperature for a given location.",
"parameters": {
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city name, e.g. San Francisco",
},
},
"required": ["location"],
},
}
}
message = [
# ESSENTIAL SYSTEM PROMPT:
# This line activates the model's function calling logic.
{
"role": "developer",
"content": "You are a model that can do function calling with the following functions"
},
{
"role": "user",
"content": "What's the temperature in London?"
}
]
inputs = processor.apply_chat_template(message, tools=[weather_function_schema], add_generation_prompt=True, return_dict=True, return_tensors="pt")
out = model.generate(**inputs.to(model.device), pad_token_id=processor.eos_token_id, max_new_tokens=128)
output = processor.decode(out[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
print(output)
<start_function_call>call:get_current_temperature{location:London}<end_function_call>
What problem are you trying to solve with this feature?
No response
Any other information you'd like to share?
No response