Coverage for transformer_lens/BertNextSentencePrediction.py: 96%

1"""Next Sentence Prediction. 

2 

3Contains a BERT style model specifically for Next Sentence Prediction. This is separate from 

4:class:`transformer_lens.HookedTransformer` because it has a significantly different architecture 

5to e.g. GPT style transformers. 

6""" 

7 

8import warnings 

9from typing import Any, Dict, List, Optional, Tuple, Union, overload 

10 

11import torch 

12from jaxtyping import Float, Int 

13from typing_extensions import Literal 

14 

15from transformer_lens.ActivationCache import ActivationCache 

16 

17 

18class BertNextSentencePrediction: 

19 """A BERT-style model for Next Sentence Prediction (NSP) that extends HookedEncoder. 

20 

21 This class implements a BERT model specifically designed for the Next Sentence Prediction task, 

22 where the model predicts whether two input sentences naturally follow each other in the original text. 

23 It inherits from HookedEncoder and adds NSP-specific components like the NSP head and pooler layer. 

24 

25 The model processes pairs of sentences and outputs either logits or human-readable predictions 

26 indicating whether the sentences are sequential. String inputs are automatically tokenized with 

27 appropriate token type IDs to distinguish between the two sentences. 

28 

29 Note: 

30 This model expects inputs to be provided as pairs of sentences. Single sentence inputs 

31 or inputs without proper sentence separation will raise errors. 

32 """ 

33 

34 def __init__(self, model: Any): 

35 warnings.warn( 

36 "BertNextSentencePrediction is deprecated and will be removed in 4.0. Use " 

37 "TransformerBridge.boot_transformers(...) for BERT-style models; see " 

38 "demos/BERT.ipynb.", 

39 DeprecationWarning, 

40 stacklevel=2, 

41 ) 

42 self.model = model 

43 

44 def __call__( 

45 self, 

46 input: Union[ 

47 List[str], 

48 Int[torch.Tensor, "batch pos"], 

49 ], 

50 return_type: Optional[Union[Literal["logits"], Literal["predictions"]]] = "logits", 

51 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

52 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

53 ) -> Optional[Union[Float[torch.Tensor, "batch 2"], str]]: 

54 """Makes the NextSentencePrediction instance callable. 

55 

56 This method delegates to the forward method, allowing the model to be called directly. 

57 The arguments and return types match the forward method exactly. 

58 """ 

59 return self.forward( 

60 input, 

61 return_type=return_type, 

62 token_type_ids=token_type_ids, 

63 one_zero_attention_mask=one_zero_attention_mask, 

64 ) 

65 

66 def to_tokens( 

67 self, 

68 input: List[str], 

69 move_to_device: bool = True, 

70 truncate: bool = True, 

71 ) -> Tuple[ 

72 Int[torch.Tensor, "batch pos"], 

73 Int[torch.Tensor, "batch pos"], 

74 Int[torch.Tensor, "batch pos"], 

75 ]: 

76 """Converts a string to a tensor of tokens. 

77 Taken mostly from the HookedTransformer implementation, but does not support default padding 

78 sides or prepend_bos. 

79 Args: 

80 input: List[str]]: The input to tokenize. 

81 move_to_device (bool): Whether to move the output tensor of tokens to the device the model lives on. Defaults to True 

82 truncate (bool): If the output tokens are too long, whether to truncate the output 

83 tokens to the model's max context window. Does nothing for shorter inputs. Defaults to 

84 True. 

85 """ 

86 

87 if len(input) != 2: 

88 raise ValueError( 

89 "Next sentence prediction task requires exactly two sentences, please provide a list of strings with each sentence as an element." 

90 ) 

91 

92 # We need to input the two sentences separately for NSP 

93 encodings = self.model.tokenizer( 

94 input[0], 

95 input[1], 

96 return_tensors="pt", 

97 padding=True, 

98 truncation=truncate, 

99 max_length=self.model.cfg.n_ctx if truncate else None, 

100 ) 

101 

102 tokens = encodings["input_ids"] 

103 token_type_ids = encodings["token_type_ids"] 

104 attention_mask = encodings["attention_mask"] 

105 

106 if move_to_device: 106 ↛ 111line 106 didn't jump to line 111 because the condition on line 106 was always true

107 tokens = tokens.to(self.model.cfg.device) 

108 token_type_ids = token_type_ids.to(self.model.cfg.device) 

109 attention_mask = attention_mask.to(self.model.cfg.device) 

110 

111 return tokens, token_type_ids, attention_mask 

112 

113 @overload 

114 def forward( 

115 self, 

116 input: Union[ 

117 List[str], 

118 Int[torch.Tensor, "batch pos"], 

119 ], 

120 return_type: Union[Literal["logits"], Literal["predictions"]], 

121 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

122 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

123 ) -> Union[Float[torch.Tensor, "batch 2"], str]: 

124 ... 

125 

126 @overload 

127 def forward( 

128 self, 

129 input: Union[ 

130 List[str], 

131 Int[torch.Tensor, "batch pos"], 

132 ], 

133 return_type: Literal[None], 

134 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

135 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

136 ) -> Optional[Union[Float[torch.Tensor, "batch 2"], str]]: 

137 ... 

138 

139 def forward( 

140 self, 

141 input: Union[ 

142 List[str], 

143 Int[torch.Tensor, "batch pos"], 

144 ], 

145 return_type: Optional[Union[Literal["logits"], Literal["predictions"]]] = "logits", 

146 token_type_ids: Optional[Int[torch.Tensor, "batch pos"]] = None, 

147 one_zero_attention_mask: Optional[Int[torch.Tensor, "batch pos"]] = None, 

148 ) -> Optional[Union[Float[torch.Tensor, "batch 2"], str]]: 

149 """Forward pass through the NextSentencePrediction module. Performs Next Sentence Prediction on a pair of sentences. 

150 

151 Args: 

152 input: The input to process. Can be one of: 

153 - List[str]: A list of two strings representing the two sentences NSP should be performed on 

154 - torch.Tensor: Input tokens as integers with shape (batch, position) 

155 return_type: Optional[str]: The type of output to return. Can be one of: 

156 - None: Return nothing, don't calculate logits 

157 - 'logits': Return logits tensor 

158 - 'predictions': Return human-readable predictions 

159 token_type_ids: Optional[torch.Tensor]: Binary ids indicating whether a token belongs 

160 to sequence A or B. For example, for two sentences: 

161 "[CLS] Sentence A [SEP] Sentence B [SEP]", token_type_ids would be 

162 [0, 0, ..., 0, 1, ..., 1, 1]. `0` represents tokens from Sentence A, 

163 `1` from Sentence B. If not provided, BERT assumes a single sequence input. 

164 This parameter gets inferred from the tokenizer if input is a string or list of strings. 

165 Shape is (batch_size, sequence_length). 

166 one_zero_attention_mask: Optional[torch.Tensor]: A binary mask which indicates 

167 which tokens should be attended to (1) and which should be ignored (0). 

168 Primarily used for padding variable-length sentences in a batch. 

169 For instance, in a batch with sentences of differing lengths, shorter 

170 sentences are padded with 0s on the right. If not provided, the model 

171 assumes all tokens should be attended to. 

172 This parameter gets inferred from the tokenizer if input is a string or list of strings. 

173 Shape is (batch_size, sequence_length). 

174 

175 Returns: 

176 Optional[torch.Tensor]: Depending on return_type: 

177 - None: Returns None if return_type is None 

178 - torch.Tensor: Returns logits if return_type is 'logits' (or if return_type is not explicitly provided) 

179 - Shape is (batch_size, 2) 

180 - str or List[str]: Returns string indicating if sentences are sequential if return_type is 'predictions' 

181 

182 Raises: 

183 ValueError: If using NSP task without proper input format or token_type_ids 

184 AssertionError: If using string input without a tokenizer 

185 """ 

186 

187 if isinstance(input, list): 

188 assert self.model.tokenizer is not None, "Must provide a tokenizer if input is a string" 

189 tokens, token_type_ids_from_tokenizer, attention_mask = self.to_tokens(input) 

190 

191 # If token_type_ids or attention mask are not provided, use the ones from the tokenizer 

192 token_type_ids = ( 

193 token_type_ids_from_tokenizer if token_type_ids is None else token_type_ids 

194 ) 

195 one_zero_attention_mask = ( 

196 attention_mask if one_zero_attention_mask is None else one_zero_attention_mask 

197 ) 

198 elif token_type_ids == None and isinstance(input, torch.Tensor): 

199 raise ValueError( 

200 "You are using the NSP task without specifying token_type_ids." 

201 "This means that the model will treat the input as a single sequence which will lead to incorrect results." 

202 "Please provide token_type_ids or use a string input." 

203 ) 

204 else: 

205 tokens = input 

206 

207 resid = self.model.encoder_output(tokens, token_type_ids, one_zero_attention_mask) 

208 

209 # NSP requires pooling (for more information see BertPooler) 

210 resid = self.model.pooler(resid) 

211 logits = self.model.nsp_head(resid) 

212 

213 if return_type == "predictions": 

214 logprobs = logits.log_softmax(dim=-1) 

215 predictions = [ 

216 "The sentences are sequential", 

217 "The sentences are NOT sequential", 

218 ] 

219 return predictions[logprobs.argmax(dim=-1).item()] 

220 

221 elif return_type == None: 

222 return None 

223 

224 return logits 

225 

226 @overload 

227 def run_with_cache( 

228 self, *model_args, return_cache_object: Literal[True] = True, **kwargs 

229 ) -> Tuple[Float[torch.Tensor, "batch 2"], ActivationCache,]: 

230 ... 

231 

232 @overload 

233 def run_with_cache( 

234 self, *model_args, return_cache_object: Literal[False], **kwargs 

235 ) -> Tuple[Float[torch.Tensor, "batch 2"], Dict[str, torch.Tensor],]: 

236 ... 

237 

238 def run_with_cache( 

239 self, 

240 *model_args, 

241 return_cache_object: bool = True, 

242 remove_batch_dim: bool = False, 

243 **kwargs, 

244 ) -> Tuple[Float[torch.Tensor, "batch 2"], Union[ActivationCache, Dict[str, torch.Tensor]],]: 

245 """ 

246 Wrapper around run_with_cache in HookedRootModule. If return_cache_object is True, 

247 this will return an ActivationCache object, with a bunch of useful HookedTransformer specific methods, 

248 otherwise it will return a dictionary of activations as in HookedRootModule. 

249 This function was copied directly from HookedTransformer. 

250 """ 

251 

252 # Create wrapper for forward function, such that run_with_cache uses 

253 # the forward function of this class and not HookedEncoder 

254 

255 class ForwardWrapper: 

256 def __init__(self, nsp): 

257 self.nsp = nsp 

258 self.original_forward = nsp.model.forward 

259 

260 def __enter__(self): 

261 # Store reference to wrapper function 

262 def wrapped_forward(*fargs, **fkwargs): 

263 return self.nsp.forward(*fargs, **fkwargs) 

264 

265 self.nsp.model.forward = wrapped_forward 

266 return self 

267 

268 def __exit__(self, exc_type, exc_val, exc_tb): 

269 # Restore original forward 

270 self.nsp.model.forward = self.original_forward 

271 

272 with ForwardWrapper(self): 

273 out, cache_dict = self.model.run_with_cache( 

274 *model_args, remove_batch_dim=remove_batch_dim, **kwargs 

275 ) 

276 if return_cache_object: 276 ↛ 280line 276 didn't jump to line 280 because the condition on line 276 was always true

277 cache = ActivationCache(cache_dict, self, has_batch_dim=not remove_batch_dim) 

278 return out, cache 

279 else: 

280 return out, cache_dict