Coverage for transformer_lens/model_bridge/supported_architectures/ernie4_5.py: 100%

1"""ERNIE 4.5 architecture adapter. 

2 

3Baidu's dense ERNIE 4.5 (``Ernie4_5ForCausalLM``) is a Llama-layout decoder — 

4RMSNorm + RoPE + GQA + gated MLP under identical module paths — with 

5config-gated biases (``use_bias``) and no BOS prepending. 

6""" 

7 

8from typing import Any 

9 

10from transformer_lens.model_bridge.supported_architectures.llama import ( 

11 LlamaArchitectureAdapter, 

12) 

13 

14 

15class Ernie4_5ArchitectureAdapter(LlamaArchitectureAdapter): 

16 """Architecture adapter for Ernie4_5ForCausalLM models.""" 

17 

18 def __init__(self, cfg: Any) -> None: 

19 super().__init__(cfg) 

20 # Verified against baidu/ERNIE-4.5-0.3B-PT's tokenizer. 

21 self.cfg.default_prepend_bos = False 

22 # ERNIE rotates adjacent element pairs (GLM-style interleaved RoPE), 

23 # unlike llama's half-split convention. 

24 self.cfg.rotary_adjacent_pairs = True 

25 # Biases are config-gated (use_bias); reshape them so a use_bias=True 

26 # GQA checkpoint gets the (n_kv, d_head) K/V bias layout. No-op when 

27 # the checkpoint carries no attention biases. 

28 self.weight_processing_conversions = { 

29 **self._qkvo_weight_conversions(include_biases=True), 

30 }