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

1"""Jais 2 architecture adapter. 

2 

3G42/Inception's Jais 2 Arabic-English family (``Jais2ForCausalLM``, 

4native in transformers): a pre-LayerNorm rotary decoder with an ungated 

5biased up/down MLP — the exact Nemotron block shape under the same module 

6names, so this subclasses Nemotron and only re-adds the attention biases. 

7""" 

8 

9from typing import Any 

10 

11from transformer_lens.model_bridge.supported_architectures.nemotron import ( 

12 NemotronArchitectureAdapter, 

13) 

14 

15 

16class Jais2ArchitectureAdapter(NemotronArchitectureAdapter): 

17 """Architecture adapter for Jais2ForCausalLM models.""" 

18 

19 # Plain nn.LayerNorm, not Nemotron's LayerNorm1P: no weight offset, foldable. 

20 _layernorm_weight_offset = False 

21 

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

23 super().__init__(cfg) 

24 # Jais 2 sets attention_bias=True; the Nemotron parent is bias-free by 

25 # default and omits bias reshapes, so Q/K/V biases would keep the flat 

26 # (n*d_head,) layout instead of (n, d_head). 

27 self.weight_processing_conversions = { 

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

29 }