Merge pull request #44 from andyrdt/andyrdt/bos_bugfix

safer remove_bos logic

Author: adamkarvonen

dictionary_learning/buffer.py

@@ -56,7 +56,7 @@ def __init__(self,
         self.refresh_batch_size = refresh_batch_size
         self.out_batch_size = out_batch_size
         self.device = device
-        self.remove_bos = remove_bos
+        self.remove_bos = remove_bos and (self.model.tokenizer.bos_token_id is not None)
         self.add_special_tokens = add_special_tokens
 
     def __iter__(self):
@@ -133,14 +133,15 @@ def refresh(self):
                     input = self.model.inputs.save()
 
                     self.submodule.output.stop()
-            attn_mask = input.value[1]["attention_mask"]
+
+            mask = (input.value[1]["attention_mask"] != 0)
             hidden_states = hidden_states.value
             if isinstance(hidden_states, tuple):
                 hidden_states = hidden_states[0]
             if self.remove_bos:
-                hidden_states = hidden_states[:, 1:, :]
-                attn_mask = attn_mask[:, 1:]
-            hidden_states = hidden_states[attn_mask != 0]
+                bos_mask = (input.value[1]["input_ids"] == self.model.tokenizer.bos_token_id)
+                mask = mask & ~bos_mask
+            hidden_states = hidden_states[mask]
 
             remaining_space = self.activation_buffer_size - current_idx
             assert remaining_space > 0

dictionary_learning/pytorch_buffer.py

@@ -117,9 +117,9 @@ def __init__(
         self.refresh_batch_size = refresh_batch_size
         self.out_batch_size = out_batch_size
         self.device = device
-        self.remove_bos = remove_bos
         self.add_special_tokens = add_special_tokens
         self.tokenizer = AutoTokenizer.from_pretrained(model.name_or_path)
+        self.remove_bos = remove_bos and (self.tokenizer.bos_token_id is not None)
 
         if not self.tokenizer.pad_token:
             self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -192,11 +192,11 @@ def refresh(self):
             with t.no_grad():
                 input = self.tokenized_batch()
                 hidden_states = collect_activations(self.model, self.submodule, input)
-            attn_mask = input["attention_mask"]
+            mask = (input["attention_mask"] != 0)
             if self.remove_bos:
-                hidden_states = hidden_states[:, 1:, :]
-                attn_mask = attn_mask[:, 1:]
-            hidden_states = hidden_states[attn_mask != 0]
+                bos_mask = (input["input_ids"] == self.tokenizer.bos_token_id)
+                mask = mask & ~bos_mask
+            hidden_states = hidden_states[mask]
 
             remaining_space = self.activation_buffer_size - current_idx
             assert remaining_space > 0