assert right padding for remove_bos logic

Author: andyrdt

dictionary_learning/buffer.py

@@ -59,6 +59,11 @@ def __init__(self,
         self.remove_bos = remove_bos
         self.add_special_tokens = add_special_tokens
 
+        print(self.model.tokenizer.padding_side)
+
+        if self.remove_bos:
+            assert self.model.tokenizer.padding_side == "right", "Padding side must be right (bos-trimming logic assumes right padding)"
+
     def __iter__(self):
         return self
 
@@ -138,6 +143,7 @@ def refresh(self):
             if isinstance(hidden_states, tuple):
                 hidden_states = hidden_states[0]
             if self.remove_bos:
+                assert self.model.tokenizer.padding_side == "right", "Padding side must be right (bos-trimming logic assumes right padding)"
                 hidden_states = hidden_states[:, 1:, :]
                 attn_mask = attn_mask[:, 1:]
             hidden_states = hidden_states[attn_mask != 0]

dictionary_learning/pytorch_buffer.py

@@ -123,6 +123,9 @@ def __init__(
 
         if not self.tokenizer.pad_token:
             self.tokenizer.pad_token = self.tokenizer.eos_token
+        
+        if self.remove_bos:
+            assert self.tokenizer.padding_side == "right", "Padding side must be right (bos-trimming logic assumes right padding)"
 
     def __iter__(self):
         return self
@@ -194,6 +197,7 @@ def refresh(self):
                 hidden_states = collect_activations(self.model, self.submodule, input)
             attn_mask = input["attention_mask"]
             if self.remove_bos:
+                assert self.tokenizer.padding_side == "right", "Padding side must be right (bos-trimming logic assumes right padding)"
                 hidden_states = hidden_states[:, 1:, :]
                 attn_mask = attn_mask[:, 1:]
             hidden_states = hidden_states[attn_mask != 0]