Coverage for transformer_lens/utilities/exploratory_utils.py: 75%

1"""attribute_utils. 

2 

3This module contains utility functions related to exploratory analysis 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import Optional, Union 

9 

10import torch 

11from rich import print as rprint 

12 

13 

14def test_prompt( 

15 prompt: str, 

16 answer: Union[str, list[str]], 

17 model, # Can't give type hint due to circular imports 

18 prepend_space_to_answer: bool = True, 

19 print_details: bool = True, 

20 prepend_bos: Optional[bool] = None, 

21 top_k: int = 10, 

22) -> None: 

23 """Test if the Model Can Give the Correct Answer to a Prompt. 

24 

25 Intended for exploratory analysis. Prints out the performance on the answer (rank, logit, prob), 

26 as well as the top k tokens. Works for multi-token prompts and multi-token answers. 

27 

28 Warning: 

29 

30 This will print the results (it does not return them). 

31 

32 Examples: 

33 

34 >>> from transformer_lens import utilities 

35 >>> from transformer_lens.model_bridge import TransformerBridge 

36 >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") 

37 >>> model.enable_compatibility_mode() 

38 

39 >>> prompt = "Why did the elephant cross the" 

40 >>> answer = "road" 

41 >>> utilities.test_prompt(prompt, answer, model) 

42 Tokenized prompt: ['<|endoftext|>', 'Why', ' did', ' the', ' elephant', ' cross', ' the'] 

43 Tokenized answer: [' road'] 

44 Performance on answer token: 

45 Rank: 2 Logit: 14.24 Prob: 3.51% Token: | road| 

46 Top 0th token. Logit: 14.51 Prob: 4.59% Token: | ground| 

47 Top 1th token. Logit: 14.41 Prob: 4.18% Token: | tree| 

48 Top 2th token. Logit: 14.24 Prob: 3.51% Token: | road| 

49 Top 3th token. Logit: 14.22 Prob: 3.45% Token: | car| 

50 Top 4th token. Logit: 13.92 Prob: 2.55% Token: | river| 

51 Top 5th token. Logit: 13.79 Prob: 2.25% Token: | street| 

52 Top 6th token. Logit: 13.77 Prob: 2.21% Token: | k| 

53 Top 7th token. Logit: 13.75 Prob: 2.16% Token: | hill| 

54 Top 8th token. Logit: 13.64 Prob: 1.92% Token: | swing| 

55 Top 9th token. Logit: 13.46 Prob: 1.61% Token: | park| 

56 Ranks of the answer tokens: [(' road', 2)] 

57 

58 Args: 

59 prompt: 

60 The prompt string, e.g. "Why did the elephant cross the". 

61 answer: 

62 The answer, e.g. "road". Note that if you set prepend_space_to_answer to False, you need 

63 to think about if you have a space before the answer here (as e.g. in this example the 

64 answer may really be " road" if the prompt ends without a trailing space). If this is a 

65 list of strings, then we only look at the next-token completion, and we compare them all 

66 as possible model answers. 

67 model: 

68 The model. 

69 prepend_space_to_answer: 

70 Whether or not to prepend a space to the answer. Note this will only ever prepend a 

71 space if the answer doesn't already start with one. 

72 print_details: 

73 Print the prompt (as a string but broken up by token), answer and top k tokens (all 

74 with logit, rank and probability). 

75 prepend_bos: 

76 Overrides self.cfg.default_prepend_bos if set. Whether to prepend 

77 the BOS token to the input (applicable when input is a string). Models generally learn 

78 to use the BOS token as a resting place for attention heads (i.e. a way for them to be 

79 "turned off"). This therefore often improves performance slightly. 

80 top_k: 

81 Top k tokens to print details of (when print_details is set to True). 

82 

83 Returns: 

84 None (just prints the results directly). 

85 """ 

86 answers = [answer] if isinstance(answer, str) else answer 

87 n_answers = len(answers) 

88 using_multiple_answers = n_answers > 1 

89 if prepend_space_to_answer: 89 ↛ 92line 89 didn't jump to line 92 because the condition on line 89 was always true

90 answers = [answer if answer.startswith(" ") else " " + answer for answer in answers] 

91 # GPT-2 often treats the first token weirdly, so lets give it a resting position 

92 prompt_tokens = model.to_tokens(prompt, prepend_bos=prepend_bos) 

93 # If we have multiple answers, we're only allowed a single token generation 

94 if using_multiple_answers: 

95 answer_tokens = torch.cat( 

96 [model.to_tokens(answer, prepend_bos=False)[:, :1] for answer in answers] 

97 ) 

98 else: 

99 answer_tokens = model.to_tokens(answers, prepend_bos=False) 

100 # Deal with case where answers is a list of strings 

101 prompt_tokens = prompt_tokens.repeat(answer_tokens.shape[0], 1) 

102 tokens = torch.cat((prompt_tokens, answer_tokens), dim=1) 

103 prompt_str_tokens = model.to_str_tokens(prompt, prepend_bos=prepend_bos) 

104 answer_str_tokens_list = [model.to_str_tokens(answer, prepend_bos=False) for answer in answers] 

105 prompt_length = len(prompt_str_tokens) 

106 answer_length = 1 if using_multiple_answers else len(answer_str_tokens_list[0]) 

107 

108 if print_details: 108 ↛ 109line 108 didn't jump to line 109 because the condition on line 108 was never true

109 print("Tokenized prompt:", prompt_str_tokens) 

110 if using_multiple_answers: 

111 print("Tokenized answers:", answer_str_tokens_list) 

112 else: 

113 print("Tokenized answer:", answer_str_tokens_list[0]) 

114 logits = model(tokens) 

115 probs = logits.softmax(dim=-1) 

116 answer_ranks = [] 

117 

118 for index in range(prompt_length, prompt_length + answer_length): 

119 # Get answer tokens for this sequence position 

120 answer_tokens = tokens[:, index] 

121 answer_str_tokens = [a[index - prompt_length] for a in answer_str_tokens_list] 

122 # Offset by 1 because models predict the NEXT token 

123 token_probs = probs[:, index - 1] 

124 sorted_token_probs, sorted_token_positions = token_probs.sort(descending=True) 

125 answer_token_ranks = sorted_token_positions.argsort(-1)[ 

126 range(n_answers), answer_tokens.cpu() 

127 ].tolist() 

128 answer_ranks.append( 

129 [ 

130 (answer_str_token, answer_token_rank) 

131 for answer_str_token, answer_token_rank in zip( 

132 answer_str_tokens, answer_token_ranks 

133 ) 

134 ] 

135 ) 

136 if print_details: 136 ↛ 139line 136 didn't jump to line 139 because the condition on line 136 was never true

137 # String formatting syntax - the first number gives the number of characters to pad to, the second number gives the number of decimal places. 

138 # rprint gives rich text printing 

139 rprint( 

140 f"Performance on answer token{'s' if n_answers > 1 else ''}:\n" 

141 + "\n".join( 

142 [ 

143 f"[b]Rank: {answer_token_ranks[i]: <8} Logit: {logits[i, index-1, answer_tokens[i]].item():5.2f} Prob: {token_probs[i, answer_tokens[i]].item():6.2%} Token: |{answer_str_tokens[i]}|[/b]" 

144 for i in range(n_answers) 

145 ] 

146 ) 

147 ) 

148 for i in range(top_k): 

149 print( 

150 f"Top {i}th token. Logit: {logits[0, index-1, sorted_token_positions[0, i]].item():5.2f} Prob: {sorted_token_probs[0, i].item():6.2%} Token: |{model.to_string(sorted_token_positions[0, i])}|" 

151 ) 

152 # If n_answers = 1 then unwrap answer ranks, so printed output matches original version of function 

153 if not using_multiple_answers: 

154 single_answer_ranks = [r[0] for r in answer_ranks] 

155 rprint(f"[b]Ranks of the answer tokens:[/b] {single_answer_ranks}") 

156 else: 

157 rprint(f"[b]Ranks of the answer tokens:[/b] {answer_ranks}") 

158 

159 

160try: 

161 import pytest 

162 

163 # Note: PyTest collects and runs this docstring as a doctest. The skip marker only applies to the 

164 # accidentally collected function item (because its name is prefixed `test_`). 

165 pytest.mark.skip(test_prompt) 

166except ModuleNotFoundError: 

167 pass # disregard if pytest not in env