Coverage for pattern_lens\load_activations.py: 85%

1"loading activations from .npz on disk. implements some custom Exception classes" 

2 

3import base64 

4import hashlib 

5import json 

6from pathlib import Path 

7from typing import Literal, overload 

8 

9import numpy as np 

10 

11 

12class GetActivationsError(ValueError): 

13 """base class for errors in getting activations""" 

14 

15 pass 

16 

17 

18class ActivationsMissingError(GetActivationsError, FileNotFoundError): 

19 """error for missing activations -- can't find the activations file""" 

20 

21 pass 

22 

23 

24class ActivationsMismatchError(GetActivationsError): 

25 """error for mismatched activations -- the prompt text or hash do not match 

26 

27 raised by `compare_prompt_to_loaded` 

28 """ 

29 

30 pass 

31 

32 

33def compare_prompt_to_loaded(prompt: dict, prompt_loaded: dict) -> None: 

34 """compare a prompt to a loaded prompt, raise an error if they do not match 

35 

36 # Parameters: 

37 - `prompt : dict` 

38 - `prompt_loaded : dict` 

39 

40 # Returns: 

41 - `None` 

42 

43 # Raises: 

44 - `ActivationsMismatchError` : if the prompt text or hash do not match 

45 """ 

46 for key in ("text", "hash"): 

47 if prompt[key] != prompt_loaded[key]: 

48 raise ActivationsMismatchError( 

49 f"Prompt file does not match prompt at key {key}:\n{prompt}\n{prompt_loaded}" 

50 ) 

51 

52 

53def augment_prompt_with_hash(prompt: dict) -> dict: 

54 """if a prompt does not have a hash, add one 

55 

56 # Parameters: 

57 - `prompt : dict` 

58 

59 # Returns: 

60 - `dict` 

61 

62 # Modifies: 

63 the input `prompt` dictionary, if it does not have a `"hash"` key 

64 """ 

65 if "hash" not in prompt: 

66 prompt_str: str = prompt["text"] 

67 prompt_hash: str = ( 

68 base64.urlsafe_b64encode(hashlib.md5(prompt_str.encode()).digest()) 

69 .decode() 

70 .rstrip("=") 

71 ) 

72 prompt.update(hash=prompt_hash) 

73 return prompt 

74 

75 

76@overload 

77def load_activations( 

78 model_name: str, 

79 prompt: dict, 

80 save_path: Path, 

81 return_fmt: Literal["torch"] = "torch", 

82) -> "tuple[Path, dict[str, torch.Tensor]]": # type: ignore[name-defined] # noqa: F821 

83 ... 

84@overload 

85def load_activations( 

86 model_name: str, 

87 prompt: dict, 

88 save_path: Path, 

89 return_fmt: Literal["numpy"] = "numpy", 

90) -> "tuple[Path, dict[str, np.ndarray]]": ... 

91def load_activations( 

92 model_name: str, 

93 prompt: dict, 

94 save_path: Path, 

95 return_fmt: Literal["torch", "numpy"] = "torch", 

96) -> "tuple[Path, dict[str, torch.Tensor]|dict[str, np.ndarray]]": # type: ignore[name-defined] # noqa: F821 

97 """load activations for a prompt and model, from an npz file 

98 

99 # Parameters: 

100 - `model_name : str` 

101 - `prompt : dict` 

102 - `save_path : Path` 

103 - `return_fmt : Literal["torch", "numpy"]` 

104 (defaults to `"torch"`) 

105 

106 # Returns: 

107 - `tuple[Path, dict[str, torch.Tensor]|dict[str, np.ndarray]]` 

108 the path to the activations file and the activations as a dictionary of numpy arrays or torch tensors, depending on `return_fmt` 

109 

110 # Raises: 

111 - `ActivationsMissingError` : if the activations file is missing 

112 - `ValueError` : if `return_fmt` is not `"torch"` or `"numpy"` 

113 """ 

114 

115 if return_fmt not in ("torch", "numpy"): 

116 raise ValueError( 

117 f"Invalid return_fmt: {return_fmt}, expected 'torch' or 'numpy'" 

118 ) 

119 if return_fmt == "torch": 

120 import torch 

121 

122 augment_prompt_with_hash(prompt) 

123 

124 prompt_dir: Path = save_path / model_name / "prompts" / prompt["hash"] 

125 prompt_file: Path = prompt_dir / "prompt.json" 

126 if not prompt_file.exists(): 

127 raise ActivationsMissingError(f"Prompt file {prompt_file} does not exist") 

128 with open(prompt_dir / "prompt.json", "r") as f: 

129 prompt_loaded: dict = json.load(f) 

130 compare_prompt_to_loaded(prompt, prompt_loaded) 

131 

132 activations_path: Path = prompt_dir / "activations.npz" 

133 

134 cache: dict 

135 

136 with np.load(activations_path) as npz_data: 

137 if return_fmt == "numpy": 

138 cache = {k: v for k, v in npz_data.items()} 

139 elif return_fmt == "torch": 

140 cache = {k: torch.from_numpy(v) for k, v in npz_data.items()} 

141 

142 return activations_path, cache