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

1"""LongT5 architecture adapter. 

2 

3Google's LongT5 (``LongT5ForConditionalGeneration``): a T5 stack whose 

4encoder self-attention is replaced by local windowed attention or 

5transient-global attention (``encoder_attention_type``). The decoder is 

6identical to T5. Encoder attention stays delegated to HF — its block-wise 

7position bias has a [1, 1, heads, block, 3*block] shape the generic 

8reconstruction cannot supply. 

9""" 

10 

11from typing import Any 

12 

13from transformer_lens.model_bridge.generalized_components import ( 

14 AttentionBridge, 

15 LinearBridge, 

16 PosEmbedBridge, 

17 RMSNormalizationBridge, 

18 T5BlockBridge, 

19) 

20from transformer_lens.model_bridge.supported_architectures.t5 import ( 

21 T5ArchitectureAdapter, 

22) 

23 

24_ENCODER_ATTN_ATTR = { 

25 "local": "LocalSelfAttention", 

26 "transient-global": "TransientGlobalSelfAttention", 

27} 

28 

29 

30class LongT5ArchitectureAdapter(T5ArchitectureAdapter): 

31 """Architecture adapter for LongT5ForConditionalGeneration models.""" 

32 

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

34 """Initialize the LongT5 architecture adapter.""" 

35 super().__init__(cfg) 

36 

37 attn_attr = _ENCODER_ATTN_ATTR[getattr(cfg, "encoder_attention_type", "local")] 

38 

39 encoder_mlp = self._build_ff_bridge("layer.1") 

40 

41 # Rebuild the encoder stack around the local/tglobal attention module; 

42 # the decoder mapping inherited from T5 is unchanged. 

43 self.components["pos_embed"] = PosEmbedBridge( 

44 name=f"encoder.block.0.layer.0.{attn_attr}.relative_attention_bias" 

45 ) 

46 self.components["encoder_blocks"] = T5BlockBridge( 

47 name="encoder.block", 

48 config=self.cfg, 

49 is_decoder=False, 

50 submodules={ 

51 "ln1": RMSNormalizationBridge(name="layer.0.layer_norm", config=self.cfg), 

52 "attn": AttentionBridge( 

53 name=f"layer.0.{attn_attr}", 

54 config=self.cfg, 

55 submodules={ 

56 "q": LinearBridge(name="q"), 

57 "k": LinearBridge(name="k"), 

58 "v": LinearBridge(name="v"), 

59 "o": LinearBridge(name="o"), 

60 }, 

61 maintain_native_attention=True, 

62 ), 

63 "ln2": RMSNormalizationBridge(name="layer.1.layer_norm", config=self.cfg), 

64 "mlp": encoder_mlp, 

65 }, 

66 )