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

1"""LED (Longformer Encoder-Decoder) architecture adapter. 

2 

3AllenAI's LED (``LEDForConditionalGeneration``: led-base/large-16384): a 

4BART-layout post-LN encoder-decoder whose encoder self-attention is 

5Longformer's sliding-window + global attention (separate query/key/value 

6and *_global projections behind an ``output`` projection). The encoder 

7attention stays delegated to HF; the decoder is plain BART attention. The 

8whole stack lives under the ``led.`` prefix instead of ``model.``. 

9 

10Encoder caveats: HF pads inputs to a multiple of config.attention_window 

11inside LEDEncoder.forward, so encoder-block hooks fire on window-padded 

12sequence lengths (only the final hidden state is unpadded). hook_q/k/v 

13cover the sliding-window projections; the global path is hookable at 

14q_global/k_global/v_global when global attention is requested. 

15""" 

16 

17from typing import Any 

18 

19from transformer_lens.model_bridge.generalized_components import ( 

20 AttentionBridge, 

21 LinearBridge, 

22) 

23from transformer_lens.model_bridge.generalized_components.base import ( 

24 CloneOutputUnderGradMixin, 

25) 

26from transformer_lens.model_bridge.supported_architectures.bart import ( 

27 BartArchitectureAdapter, 

28) 

29 

30 

31class _LEDEncoderQueryBridge(CloneOutputUnderGradMixin, LinearBridge): 

32 """LEDEncoderSelfAttention scales the query projection with an in-place 

33 ``/=``; clone under grad (see mixin).""" 

34 

35 

36class LEDArchitectureAdapter(BartArchitectureAdapter): 

37 """Architecture adapter for LEDForConditionalGeneration models.""" 

38 

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

40 """Initialize the LED architecture adapter.""" 

41 super().__init__(cfg) 

42 

43 self._reprefix_components("model.", "led.") 

44 

45 def _encoder_attention(self) -> AttentionBridge: 

46 """Sliding-window + global attention; window chunking and the global 

47 projections have no generic reconstruction, so the module delegates.""" 

48 return AttentionBridge( 

49 name="self_attn", 

50 config=self.cfg, 

51 submodules={ 

52 "q": _LEDEncoderQueryBridge(name="longformer_self_attn.query"), 

53 "k": LinearBridge(name="longformer_self_attn.key"), 

54 "v": LinearBridge(name="longformer_self_attn.value"), 

55 "q_global": LinearBridge(name="longformer_self_attn.query_global"), 

56 "k_global": LinearBridge(name="longformer_self_attn.key_global"), 

57 "v_global": LinearBridge(name="longformer_self_attn.value_global"), 

58 "o": LinearBridge(name="output"), 

59 }, 

60 maintain_native_attention=True, 

61 )