Coverage for transformer_lens/model_bridge/generalized_components/pos_embed.py: 90%

1"""Positional embedding bridge component. 

2 

3This module contains the bridge component for positional embedding layers. 

4""" 

5from typing import Any, Dict, Optional 

6 

7import torch 

8 

9from transformer_lens.model_bridge.generalized_components.base import ( 

10 GeneralizedComponent, 

11) 

12 

13 

14class PosEmbedBridge(GeneralizedComponent): 

15 """Positional embedding bridge that wraps transformer positional embedding layers. 

16 

17 This component provides standardized input/output hooks for positional embeddings. 

18 """ 

19 

20 property_aliases = {"W_pos": "weight"} 

21 

22 def __init__( 

23 self, 

24 name: str, 

25 config: Optional[Any] = None, 

26 submodules: Optional[Dict[str, GeneralizedComponent]] = {}, 

27 ): 

28 """Initialize the positional embedding bridge. 

29 

30 Args: 

31 name: The name of this component 

32 config: Optional configuration (unused for PosEmbedBridge) 

33 submodules: Dictionary of GeneralizedComponent submodules to register 

34 """ 

35 super().__init__(name, config, submodules=submodules) 

36 

37 @property 

38 def W_pos(self) -> torch.Tensor: 

39 """Return the positional embedding weight matrix.""" 

40 if self.original_component is None: 40 ↛ 41line 40 didn't jump to line 41 because the condition on line 40 was never true

41 raise RuntimeError(f"Original component not set for {self.name}") 

42 # Sinusoidal nn.Module variants (e.g. M2M100) register the table as a 

43 # "weights" buffer instead of an nn.Embedding "weight" parameter. 

44 weight = getattr(self.original_component, "weight", None) 

45 if weight is None: 

46 weight = getattr(self.original_component, "weights", None) 

47 assert isinstance( 

48 weight, torch.Tensor 

49 ), f"Component {self.name} has no weight/weights tensor" 

50 return weight 

51 

52 def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor: 

53 """Forward pass through the positional embedding bridge. 

54 

55 This method accepts variable arguments to support different architectures: 

56 - Standard models (GPT-2, GPT-Neo): (input_ids, position_ids=None) 

57 - OPT models: (attention_mask, past_key_values_length=0, position_ids=None) 

58 - Others may have different signatures 

59 

60 Args: 

61 *args: Positional arguments forwarded to the original component 

62 **kwargs: Keyword arguments forwarded to the original component 

63 

64 Returns: 

65 Positional embeddings 

66 """ 

67 if self.original_component is None: 67 ↛ 68line 67 didn't jump to line 68 because the condition on line 67 was never true

68 raise RuntimeError( 

69 f"Original component not set for {self.name}. Call set_original_component() first." 

70 ) 

71 # Sinusoidal variants (e.g. Marian) receive a torch.Size, not a tensor; 

72 # there is nothing to hook in that case. 

73 if args and isinstance(args[0], torch.Tensor): 

74 first_arg = self.hook_in(args[0]) 

75 args = (first_arg,) + args[1:] 

76 output = self.original_component(*args, **kwargs) 

77 

78 # Expand batch=1 pos embeddings to match actual batch size for hooks. 

79 batch_size = getattr(self, "_current_batch_size", None) 

80 

81 # Read-and-clear to avoid stale values during generate() steps. 

82 if batch_size is not None: 

83 self._current_batch_size = None 

84 if ( 

85 batch_size is not None 

86 and batch_size > 1 

87 and isinstance(output, torch.Tensor) 

88 and output.ndim >= 1 

89 and output.shape[0] == 1 

90 ): 

91 output = output.expand(batch_size, *[-1] * (output.ndim - 1)) 

92 

93 output = self.hook_out(output) 

94 return output