Skip to content

Add export file for FunctionGemma such as export to tflite or onnx, and the test file #272

@AFun3

Description

@AFun3

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions