maxtext.experimental.agent.ckpt_conversion_agent.mapping module
A mapping agent, to generate param_mappings and hf_shape
-
class maxtext.experimental.agent.ckpt_conversion_agent.mapping.MappingAgent(api_key, dir_path, target_model='gemma3-4b', max_retries=3)[source]
Bases: BaseAgent
An agent that generates and verifies mapping functions for model conversion.
-
generate_param_mapping()[source]
Generates and verifies the parameter mapping function.
-
generate_shape_mapping()[source]
Generates and verifies the Hugging Face weights shape mapping function.
