Function transformers_gemma3_from_encoder

gemma_encoder.py:588–671

(encoder)

def transformers_gemma3_from_encoder(encoder):
    jsons_path = Path(__file__).parent / "gemma_configs"
    config = Gemma3Config.from_json_file(jsons_path / "gemma3cfg.json")
    with torch.device("meta"):
        metamodel = Gemma3ForConditionalGeneration(config)
    t_model_name = config.text_config.model_type
    t_model = _locate_model_within_model(metamodel, t_model_name)
    if t_model is None:
        raise ValueError(
            "Can't locate text model while converting text encoder to Gemma3ForConditionalGeneration"
        )
    t_model.load_state_dict(
        encoder.gemma3_12b.transformer.model.state_dict(), assign=True, strict=False
    )
    v_tower_name = config.vision_config.model_type
    v_tower = _locate_model_within_model(metamodel, v_tower_name)
    if v_tower is None:
        raise ValueError(
            "Can't locate vision model while converting text encoder to Gemma3ForConditionalGeneration"
        )
    v_model = v_tower.vision_model
    v_model.load_state_dict(
        encoder.gemma3_12b.transformer.vision_model.state_dict(),
        assign=True,
        strict=False,
    )
    metamodel.multi_modal_projector.load_state_dict(
        encoder.gemma3_12b.transformer.multi_modal_projector.state_dict(),
        assign=True,
        strict=False,
    )
    config = config.text_config
    dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
    base = config.rope_local_base_freq

    device = encoder.device
    positions_length = len(v_model.embeddings.position_ids[0])
    position_ids = torch.arange(
        positions_length, dtype=torch.long, device="cpu"
    ).unsqueeze(0)
    v_model.embeddings.register_buffer("position_ids", position_ids)
    embed_scale = torch.tensor(config.hidden_size**0.5, device=device)
    t_model.embed_tokens.register_buffer("embed_scale", embed_scale)
    local_rope_freqs = 1.0 / (
        base
        ** (
            torch.arange(0, dim, 2, dtype=torch.int64).to(
                device=device, dtype=torch.float
            )
            / dim
        )
    )
    t_model.rotary_emb_local.register_buffer("inv_freq", local_rope_freqs)
    rope_freqs, _ = ROPE_INIT_FUNCTIONS[config.rope_scaling["rope_type"]](
        config, device
    )
    t_model.rotary_emb.register_buffer("inv_freq", rope_freqs)
    lm_head_requires_grad = False