Coverage for transformer_lens/conversion_utils/conversion_steps/tensor_conversion_set.py: 61%

1"""Tensor conversion set.""" 

2 

3from typing import Any 

4 

5import torch 

6 

7from transformer_lens.conversion_utils.helpers.find_property import find_property 

8from transformer_lens.conversion_utils.hook_conversion_utils import ( 

9 get_weight_conversion_field_set, 

10) 

11 

12from .base_tensor_conversion import BaseTensorConversion 

13from .rearrange_tensor_conversion import RearrangeTensorConversion 

14 

15 

16class TensorConversionSet(BaseTensorConversion): 

17 def __init__( 

18 self, 

19 fields: dict[str, Any], 

20 ): 

21 super().__init__() 

22 self.fields = fields 

23 

24 def get_component(self, model: Any, name: str) -> Any: 

25 """Get a component from the model using the field mapping. 

26 

27 Args: 

28 model: The model to get the component from. 

29 name: The name of the component to get. 

30 

31 Returns: 

32 The requested component. 

33 """ 

34 if name not in self.fields: 

35 raise ValueError(f"Unknown component name: {name}") 

36 

37 field_info = self.fields[name] 

38 if isinstance(field_info, str): 

39 field_name = field_info 

40 conversion_step = None 

41 else: 

42 field_name, conversion_step = field_info 

43 

44 component = find_property(field_name, model) 

45 

46 if conversion_step is not None: 

47 component = conversion_step(component) 

48 

49 return component 

50 

51 def handle_conversion(self, input_value: Any, *full_context: Any) -> dict[str, Any]: 

52 result = {} 

53 for fields_name in self.fields: 

54 conversion_action = self.fields[fields_name] 

55 result[fields_name] = self.process_conversion_action( 

56 input_value, 

57 conversion_details=conversion_action, 

58 ) 

59 

60 return result 

61 

62 def process_conversion_action( 

63 self, input_value: Any, conversion_details: Any, *full_context: Any 

64 ) -> Any: 

65 if isinstance(conversion_details, torch.Tensor): 

66 return conversion_details 

67 elif isinstance(conversion_details, str): 

68 return find_property(conversion_details, input_value) 

69 else: 

70 (remote_field, conversion) = conversion_details 

71 return self.process_conversion(input_value, remote_field, conversion, *full_context) 

72 

73 def process_conversion( 

74 self, 

75 input_value: Any, 

76 remote_field: str, 

77 conversion: BaseTensorConversion, 

78 *full_context: Any, 

79 ) -> Any: 

80 field = find_property(remote_field, input_value) 

81 if isinstance(conversion, TensorConversionSet): 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 result = [] 

83 for layer in field: 

84 result.append(conversion.convert(layer, input_value, *full_context)) 

85 return result 

86 

87 else: 

88 return conversion.convert(field, *[input_value, *full_context]) 

89 

90 def get_conversion_action(self, field: str) -> BaseTensorConversion: 

91 conversion_details = self.fields[field] 

92 if isinstance(conversion_details, tuple): 

93 return conversion_details[1] 

94 else: 

95 # Return no op if not a specific conversion 

96 return RearrangeTensorConversion("... -> ...") 

97 

98 def __repr__(self) -> str: 

99 conversion_string = ( 

100 "Is composed of a set of nested conversions with the following details {\n\t" 

101 ) 

102 # This is a bit of a hack to get the string representation of nested conversions 

103 conversion_string += get_weight_conversion_field_set(self.fields)[:-1].replace("\n", "\n\t") 

104 conversion_string += "\n}" 

105 return conversion_string