Coverage for transformer_lens/conversion_utils/hook_conversion_utils.py: 72%

1"""Weight conversion utilities.""" 

2 

3import torch 

4 

5 

6def get_weight_conversion_field_set(weights: dict) -> str: 

7 """Creates a formatted string showing how weights are mapped between frameworks. 

8 

9 Args: 

10 weights: Dictionary containing weight mappings where: 

11 - keys are TransformerLens weight names 

12 - values can be: 

13 * tuple[str, "BaseTensorConversion"] 

14 * torch.Tensor 

15 * strings 

16 

17 Returns: 

18 A formatted multi-line string showing each weight's mapping details. 

19 """ 

20 conversion_string = "" 

21 for transformer_lens_weight in weights: 

22 hugging_face_weight = weights[transformer_lens_weight] 

23 

24 # Case 1: Nested conversion, call __repr__ for the nested conversion 

25 if isinstance(hugging_face_weight, tuple): 

26 weight_name, conversion = hugging_face_weight 

27 conversion_string += ( 

28 f'"{transformer_lens_weight}" -> "{weight_name}", {conversion.__repr__()}\n' 

29 ) 

30 

31 # Case 2: Tensor, display shape and content 

32 elif isinstance(hugging_face_weight, torch.Tensor): 

33 if torch.all(hugging_face_weight == 0): 33 ↛ 35line 33 didn't jump to line 35 because the condition on line 33 was always true

34 conversion_string += f'"{transformer_lens_weight}" -> "Tensor filled with zeros of shape {hugging_face_weight.shape}",\n' 

35 elif torch.all(hugging_face_weight == 1): 

36 conversion_string += f'"{transformer_lens_weight}" -> "Tensor filled with ones of shape {hugging_face_weight.shape}",\n' 

37 else: 

38 conversion_string += f'"{transformer_lens_weight}" -> "Tensor of shape {hugging_face_weight.shape}",\n' 

39 

40 # Case 3: String, just display string (name of weight in HuggingFace) 

41 else: 

42 conversion_string += f'"{transformer_lens_weight}" -> "{hugging_face_weight}",\n' 

43 return conversion_string 

44 

45 

46def model_info_cfg(cfg): 

47 """ 

48 Displays the weight conversion from HuggingFace to TransformerLens for a given model configuration. 

49 

50 Args: 

51 cfg: Model configuration object containing architecture information 

52 """ 

53 

54 # TODO: WeightConversionFactory import needs to be updated or removed 

55 print(f"Hook conversion details for architecture {cfg.original_architecture}:") 

56 print("Hook conversion factory not yet implemented")