| """Load the pinned ONNX tensors into the same frozen Torch backbone used in evaluation.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import onnx |
| import torch |
| from onnx import numpy_helper |
| from transformers import LlamaConfig, LlamaModel |
| from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding |
|
|
| from .core import projection_weights |
|
|
|
|
| def load_backbone(root: Path, device: str) -> LlamaModel: |
| source = json.loads((root / "config.json").read_text()) |
| graph = onnx.load(root / "emma6.onnx", load_external_data=False) |
| tensors = {} |
| for tensor in graph.graph.initializer: |
| if tensor.data_type != onnx.TensorProto.FLOAT: |
| continue |
| data = {entry.key: entry.value for entry in tensor.external_data} |
| if data: |
| array = np.memmap(root / data["location"], dtype=np.float32, mode="c", |
| offset=int(data.get("offset", 0)), shape=tuple(tensor.dims)) |
| else: |
| array = numpy_helper.to_array(tensor).copy() |
| tensors[tensor.name] = torch.from_numpy(array) |
| config = LlamaConfig( |
| vocab_size=source["vocab_size"], hidden_size=source["hidden_size"], |
| intermediate_size=source["intermediate_size"], num_hidden_layers=source["num_hidden_layers"], |
| num_attention_heads=source["num_attention_heads"], num_key_value_heads=source["num_kv_heads"], |
| max_position_embeddings=source["context_length"], rms_norm_eps=1e-6, |
| rope_parameters={"rope_type": "default", "rope_theta": source["rope_theta"]}, |
| attention_bias=False, mlp_bias=False, use_cache=False, pad_token_id=0, |
| bos_token_id=2, eos_token_id=3, |
| ) |
| config._attn_implementation = "sdpa" |
| with torch.device("meta"): |
| model = LlamaModel(config) |
| state = projection_weights(graph, tensors) |
| state.update({"embed_tokens.weight": tensors["token_emb.weight"], "norm.weight": tensors["norm_f.weight"]}) |
| for i in range(config.num_hidden_layers): |
| state[f"layers.{i}.input_layernorm.weight"] = tensors[f"blocks.{i}.norm1.weight"] |
| state[f"layers.{i}.post_attention_layernorm.weight"] = tensors[f"blocks.{i}.norm2.weight"] |
| model.load_state_dict(state, strict=True, assign=True) |
| model.rotary_emb = LlamaRotaryEmbedding(config) |
| return model.eval().requires_grad_(False).to(device) |
|
|
