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

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: 

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 answer_tokens = model.to_tokens(answers, prepend_bos=False) 

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

95 if using_multiple_answers: 

96 answer_tokens = answer_tokens[:, :1] 

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

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

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

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

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

102 prompt_length = len(prompt_str_tokens) 

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

104 

105 if print_details: 

106 print("Tokenized prompt:", prompt_str_tokens) 

107 if using_multiple_answers: 

108 print("Tokenized answers:", answer_str_tokens_list) 

109 else: 

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

111 logits = model(tokens) 

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

113 answer_ranks = [] 

114 

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

116 # Get answer tokens for this sequence position 

117 answer_tokens = tokens[:, index] 

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

119 # Offset by 1 because models predict the NEXT token 

120 token_probs = probs[:, index - 1] 

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

122 answer_token_ranks = sorted_token_positions.argsort(-1)[ 

123 range(n_answers), answer_tokens.cpu() 

124 ].tolist() 

125 answer_ranks.append( 

126 [ 

127 (answer_str_token, answer_token_rank) 

128 for answer_str_token, answer_token_rank in zip( 

129 answer_str_tokens, answer_token_ranks 

130 ) 

131 ] 

132 ) 

133 if print_details: 

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

135 # rprint gives rich text printing 

136 rprint( 

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

138 + "\n".join( 

139 [ 

140 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]" 

141 for i in range(n_answers) 

142 ] 

143 ) 

144 ) 

145 for i in range(top_k): 

146 print( 

147 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])}|" 

148 ) 

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

150 if not using_multiple_answers: 

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

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

153 else: 

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

155 

156 

157try: 

158 import pytest 

159 

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

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

162 pytest.mark.skip(test_prompt) 

163except ModuleNotFoundError: 

164 pass # disregard if pytest not in env