maxtext.experimental.agent.ckpt_conversion_agent.transformation module

maxtext.experimental.agent.ckpt_conversion_agent.transformation module#

A transformation agent to generate the layerwise and bidirectional transformation hook functions between HF & Maxtext

class maxtext.experimental.agent.ckpt_conversion_agent.transformation.TransformationAgent(api_key, dir_path, target_model='gemma3')[source]#

Bases: BaseAgent

An agent that generates transformation hook functions for model conversion.

generate_hook_functions()[source]#

Generates layer-wise transformation hook functions.