Coverage for transformer_lens/model_bridge/generalized_components/embedding.py: 66%

1"""Embedding bridge component. 

2 

3This module contains the bridge component for embedding layers. 

4""" 

5import inspect 

6from typing import Any, Dict, Optional 

7 

8import torch 

9 

10from transformer_lens.model_bridge.generalized_components.base import ( 

11 GeneralizedComponent, 

12) 

13 

14 

15class EmbeddingBridge(GeneralizedComponent): 

16 """Embedding bridge that wraps transformer embedding layers. 

17 

18 This component provides standardized input/output hooks. 

19 """ 

20 

21 property_aliases = {"W_E": "e.weight", "W_pos": "pos.weight"} 

22 

23 def __init__( 

24 self, 

25 name: str, 

26 config: Optional[Any] = None, 

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

28 output_contribution_scale: float = 1.0, 

29 ): 

30 """Initialize the embedding bridge. 

31 

32 Args: 

33 name: The name of this component 

34 config: Optional configuration (unused for EmbeddingBridge) 

35 submodules: Dictionary of GeneralizedComponent submodules to register 

36 output_contribution_scale: Outer HF runtime scale exposed by hook_out. 

37 Observation preserves the raw forward; edits are unscaled before returning. 

38 """ 

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

40 self._output_contribution_scale = float(output_contribution_scale) 

41 

42 @property 

43 def W_E(self) -> torch.Tensor: 

44 """Return the embedding weight matrix.""" 

45 if hasattr(self, "_use_processed_weights") and self._use_processed_weights: 45 ↛ 46line 45 didn't jump to line 46 because the condition on line 45 was never true

46 if hasattr(self, "_processed_weight"): 

47 return self._processed_weight 

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

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

50 if hasattr(self.original_component, "inv_freq") and ( 50 ↛ 53line 50 didn't jump to line 53 because the condition on line 50 was never true

51 not hasattr(self.original_component, "weight") 

52 ): 

53 inv_freq = self.original_component.inv_freq 

54 assert isinstance(inv_freq, torch.Tensor), f"inv_freq is not a tensor for {self.name}" 

55 return inv_freq 

56 assert hasattr( 

57 self.original_component, "weight" 

58 ), f"Component {self.name} has neither weight nor inv_freq attribute" 

59 weight = self.original_component.weight 

60 assert isinstance(weight, torch.Tensor), f"Weight is not a tensor for {self.name}" 

61 return weight 

62 

63 def forward( 

64 self, input_ids: torch.Tensor, position_ids: torch.Tensor | None = None, **kwargs: Any 

65 ) -> torch.Tensor: 

66 """Forward pass through the embedding bridge. 

67 

68 Args: 

69 input_ids: Input token IDs 

70 position_ids: Optional position IDs (ignored if not supported) 

71 **kwargs: Additional arguments to pass to the original component 

72 

73 Returns: 

74 Embedded output 

75 """ 

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

77 raise RuntimeError( 

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

79 ) 

80 target_dtype = None 

81 try: 

82 target_dtype = next(self.original_component.parameters()).dtype 

83 except StopIteration: 

84 pass 

85 input_ids = self.hook_in(input_ids) 

86 sig = inspect.signature(self.original_component.forward) 

87 supports_position_ids = "position_ids" in sig.parameters 

88 if not hasattr(self.original_component, "forward") or not supports_position_ids: 88 ↛ 92line 88 didn't jump to line 92 because the condition on line 88 was always true

89 kwargs.pop("position_ids", None) 

90 output = self.original_component(input_ids, **kwargs) 

91 else: 

92 output = self.original_component(input_ids, position_ids=position_ids, **kwargs) 

93 if isinstance(output, tuple): 93 ↛ 94line 93 didn't jump to line 94 because the condition on line 93 was never true

94 output = output[0] 

95 if target_dtype is not None and output.dtype != target_dtype: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true

96 output = output.to(dtype=target_dtype) 

97 output = self._hook_scaled_output(output, self.hook_out, self._output_contribution_scale) 

98 return output