misc

src/repe/rep_control_pipeline.py

@@ -35,7 +35,7 @@ class RepControlPipeline(RepControlPipeline2):
         self.layers = layers
         self.max_length = max_length
 
-        super().__init__(model=model, tokenizer=tokenizer, **kwargs)
+        super().__init__(model=model, tokenizer=tokenizer, max_length=max_length, **kwargs)
         
     def __call__(self, text_inputs, activations=None, **kwargs):
         if activations is not None:

src/repe/rep_control_pipeline_baukit.py

@@ -20,6 +20,7 @@ def row_choice_ids(answer_choices, tokenizer):
 
 
 def intervene(output, activation):
+    # TODO need attention mask
     assert output.ndim == 3, f"expected output to be (batch, seq, vocab), got {output.shape}"
     return output + activation.to(output.device)[None, None, :]