Coverage for transformer_lens/utilities/bridge_components.py: 85%

1"""Utilities for traversing and applying functions to every component in a TransformerBridge model.""" 

2 

3from typing import Any, Callable, cast 

4 

5import torch.nn as nn 

6 

7from transformer_lens.model_bridge.bridge import TransformerBridge 

8from transformer_lens.model_bridge.generalized_components.base import ( 

9 GeneralizedComponent, 

10) 

11 

12 

13def collect_all_submodules_of_component( 

14 model: TransformerBridge, 

15 component: GeneralizedComponent, 

16 submodules: dict, 

17 block_prefix: str = "", 

18) -> dict: 

19 """Recursively collects all submodules of a component in a TransformerBridge model. 

20 Args: 

21 model: The TransformerBridge model to collect submodules from 

22 component: The component to collect submodules from 

23 submodules: A dictionary to populate with submodules (modified in-place) 

24 block_prefix: Prefix for the block name, needed for components that are part of a block bridge 

25 Returns: 

26 Dictionary mapping submodule names to their respective submodules 

27 """ 

28 for component_submodule in component.submodules.values(): 

29 # Skip components without names (e.g., OPT's MLP container) 

30 if component_submodule.name is not None: 

31 submodules[block_prefix + component_submodule.name] = component_submodule 

32 

33 if component_submodule.is_list_item: 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 submodules = collect_components_of_block_bridge(model, component_submodule, submodules) 

35 

36 if component_submodule.submodules: 

37 submodules = collect_all_submodules_of_component( 

38 model, component_submodule, submodules, block_prefix 

39 ) 

40 return submodules 

41 

42 

43def collect_components_of_block_bridge( 

44 model: TransformerBridge, component: GeneralizedComponent, components: dict 

45) -> dict: 

46 """Collects all components of a BlockBridge component. 

47 Args: 

48 model: The TransformerBridge model to collect components from 

49 component: The BlockBridge component to collect components from 

50 components: A dictionary to populate with components (modified in-place) 

51 Returns: 

52 Dictionary mapping component names to their respective components 

53 """ 

54 

55 # Retrieve the remote component list from the adapter (we need a ModuleList to iterate over) 

56 if component.name is None: 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true

57 raise ValueError("Block bridge component must have a name") 

58 

59 # Use cached original_component for nested list items (relative names) 

60 if component.original_component is not None: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true

61 remote_module_list = component.original_component 

62 else: 

63 try: 

64 remote_module_list = model.adapter.get_remote_component( 

65 model.original_model, component.name 

66 ) 

67 except AttributeError: 

68 # Relative name not reachable from root; already set up during boot 

69 return components 

70 

71 # Make sure the remote component is a ModuleList 

72 if isinstance(remote_module_list, nn.ModuleList): 72 ↛ 81line 72 didn't jump to line 81 because the condition on line 72 was always true

73 for block in remote_module_list: 

74 block_component = cast(GeneralizedComponent, block) 

75 block_name = block_component.name 

76 assert block_name is not None, "Block bridge component must have a name" 

77 components[block_name] = block_component 

78 components = collect_all_submodules_of_component( 

79 model, block_component, components, block_name 

80 ) 

81 return components 

82 

83 

84def collect_all_components(model: TransformerBridge, components: dict) -> dict: 

85 """Collects all components in a TransformerBridge inside a dictionary. 

86 The keys are the component names, and the values are the components themselves. 

87 Args: 

88 model: The TransformerBridge model to collect components from 

89 components: A dictionary to populate with components (modified in-place) 

90 Returns: 

91 Dictionary mapping component names to their respective components 

92 """ 

93 

94 for component in model.adapter.get_component_mapping().values(): 

95 components[component.name] = component 

96 components = collect_all_submodules_of_component(model, component, components) 

97 

98 # We need to enable compatibility mode for all different blocks of the component if the component is a list item 

99 if component.is_list_item: 

100 components = collect_components_of_block_bridge(model, component, components) 

101 return components 

102 

103 

104def apply_fn_to_all_components( 

105 model: TransformerBridge, 

106 fn: Callable[[GeneralizedComponent], Any], 

107 components: dict | None = None, 

108) -> dict[str, Any]: 

109 """Applies a function to all components in the TransformerBridge model. 

110 Args: 

111 model: The TransformerBridge model to apply the function to 

112 fn: The function to apply to each component 

113 components: Optional dictionary of components to apply the function to, if None, all components are collected 

114 Returns: 

115 return_values: A dictionary mapping component names to the return values of the function 

116 """ 

117 

118 if components is None: 118 ↛ 121line 118 didn't jump to line 121 because the condition on line 118 was always true

119 components = collect_all_components(model, {}) 

120 

121 return_values = {} 

122 

123 for component in components.values(): 

124 return_values[component.name] = fn(component) 

125 

126 return return_values