Coverage for pattern_lens/load_activations.py: 89%

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 

11from pattern_lens.consts import ReturnCache 

12 

13 

14class GetActivationsError(ValueError): 

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

16 

17 pass 

18 

19 

20class ActivationsMissingError(GetActivationsError, FileNotFoundError): 

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

22 

23 pass 

24 

25 

26class ActivationsMismatchError(GetActivationsError): 

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

28 

29 raised by `compare_prompt_to_loaded` 

30 """ 

31 

32 pass 

33 

34 

35class InvalidPromptError(GetActivationsError): 

36 """error for invalid prompt -- the prompt does not have fields "hash" or "text" 

37 

38 raised by `augment_prompt_with_hash` 

39 """ 

40 

41 pass 

42 

43 

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

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

46 

47 # Parameters: 

48 - `prompt : dict` 

49 - `prompt_loaded : dict` 

50 

51 # Returns: 

52 - `None` 

53 

54 # Raises: 

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

56 """ 

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

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

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

60 raise ActivationsMismatchError( 

61 msg, 

62 ) 

63 

64 

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

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

67 

68 not having a "text" field is allowed, but only if "hash" is present 

69 

70 # Parameters: 

71 - `prompt : dict` 

72 

73 # Returns: 

74 - `dict` 

75 

76 # Modifies: 

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

78 """ 

79 if "hash" not in prompt: 

80 if "text" not in prompt: 

81 msg = f"Prompt does not have 'text' field or 'hash' field: {prompt}" 

82 raise InvalidPromptError( 

83 msg, 

84 ) 

85 prompt_str: str = prompt["text"] 

86 prompt_hash: str = ( 

87 # we don't need this to be a secure hash 

88 base64.urlsafe_b64encode(hashlib.md5(prompt_str.encode()).digest()) # noqa: S324 

89 .decode() 

90 .rstrip("=") 

91 ) 

92 prompt.update(hash=prompt_hash) 

93 return prompt 

94 

95 

96@overload 

97def load_activations( 

98 model_name: str, 

99 prompt: dict, 

100 save_path: Path, 

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

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

103 ... 

104@overload 

105def load_activations( 

106 model_name: str, 

107 prompt: dict, 

108 save_path: Path, 

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

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

111def load_activations( 

112 model_name: str, 

113 prompt: dict, 

114 save_path: Path, 

115 return_fmt: ReturnCache = "torch", 

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

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

118 

119 # Parameters: 

120 - `model_name : str` 

121 - `prompt : dict` 

122 - `save_path : Path` 

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

124 (defaults to `"torch"`) 

125 

126 # Returns: 

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

128 the path to the activations file and the activations as a dictionary 

129 of numpy arrays or torch tensors, depending on `return_fmt` 

130 

131 # Raises: 

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

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

134 """ 

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

136 msg = f"Invalid return_fmt: {return_fmt}, expected 'torch' or 'numpy'" 

137 raise ValueError( 

138 msg, 

139 ) 

140 if return_fmt == "torch": 

141 import torch 

142 

143 augment_prompt_with_hash(prompt) 

144 

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

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

147 if not prompt_file.exists(): 

148 msg = f"Prompt file {prompt_file} does not exist" 

149 raise ActivationsMissingError(msg) 

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

151 prompt_loaded: dict = json.load(f) 

152 compare_prompt_to_loaded(prompt, prompt_loaded) 

153 

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

155 

156 cache: dict 

157 

158 with np.load(activations_path) as npz_data: 

159 if return_fmt == "numpy": 

160 cache = dict(npz_data.items()) 

161 elif return_fmt == "torch": 

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

163 

164 return activations_path, cache 

165 

166 

167# def load_activations_stacked()