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 def __init__(self, cfg: Any) -> None: 

20 super().__init__(cfg) 

21 # Nemotron's fold/center disable is for LayerNorm1P (weight+1 gamma); 

22 # Jais 2 uses plain nn.LayerNorm, the standard foldable case. 

23 self.supports_fold_ln = True 

24 self.supports_center_writing_weights = True 

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

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

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

28 self.weight_processing_conversions = { 

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

30 }