Skip to content

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

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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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