misc

notebooks/make_dataset2.py

@@ -96,7 +96,7 @@ def load_rep_reader(model, tokenizer, cfg, N_fit_examples=20, batch_size=2, rep_
     intervention_f.parent.mkdir(exist_ok=True, parents=True)
     if not intervention_f.exists():        
         
-        hidden_layers = list(range(8, model.config.num_hidden_layers, 3))
+        hidden_layers = list(range(cfg.layer_padding, model.config.num_hidden_layers, cfg.layer_stride))
         
         dataset_fit = load_preproc_dataset('imdb', cfg, tokenizer, N=N_fit_examples)
         
@@ -144,7 +144,7 @@ hidden_layers
 
 # %%
 # load dataset
-ds_name = 'imdb'
+ds_name = cfg.datasets[0]
 ds_tokens = load_preproc_dataset(ds_name, tokenizer, N=sum(cfg.max_examples), seed=cfg.seed, num_shots=cfg.num_shots, max_length=cfg.max_length)
 
 N_train_split = (len(ds_tokens) - N_fit_examples) //2
@@ -153,7 +153,8 @@ N_train_split = (len(ds_tokens) - N_fit_examples) //2
 dataset_fit = ds_tokens.select(range(N_fit_examples))
 dataset_train = ds_tokens.select(range(N_fit_examples, N_train_split))
 dataset_test = ds_tokens.select(range(N_train_split, len(ds_tokens)))
-dataset_test
+assert len(dataset_train)>3, f"dataset_train is too small {len(dataset_train)}"
+assert len(dataset_test)>3
 
 
 
@@ -239,6 +240,7 @@ if TEST:
 
 # %%
 # test intervention quality
+# TODO perhaps move this to intervention create/load/cache
 if TEST:
     model.eval()
     with torch.no_grad():
@@ -295,7 +297,7 @@ def create_hs_ds(ds_name, ds_tokens, pipeline, activations=None, f = None, batch
 
 
 
-ds1, f = create_hs_ds('imdb', dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size)
+ds1, f = create_hs_ds(ds_name, dataset_train, rep_control_pipeline2, split_type="train", debug=True, batch_size=batch_size, activations=activations)
 ds1
 
 # TODO add qc

src/extraction/config.py

@@ -36,7 +36,7 @@ class ExtractConfig(Serializable):
     """Shortcut for `layers = (0,) + tuple(range(1, num_layers + 1, stride))`."""
     
     layer_padding: InitVar[int] = 4
-    """Clips the first layers by this amount"""
+    """Skips this amount of first layers"""
 
     seed: int = 42
     """Seed to use for prompt randomization. Defaults to 42."""
@@ -47,5 +47,5 @@ class ExtractConfig(Serializable):
     template_path: str | None = None
     """Path to pass into `DatasetTemplates`. By default we use the dataset name."""
     
-    max_length: int | None = 555
+    max_length: int | None = 666
     """Maximum length of the input sequence passed to the tokenize encoder function"""

src/prompts/prompt_loading.py

@@ -346,4 +346,5 @@ def load_preproc_dataset(ds_name: str, tokenizer: PreTrainedTokenizerBase, N:int
     # ## Filter out truncated examples
     ds_tokens = ds_tokens.filter(lambda r: not r['truncated'])
     print('num_rows (after filtering out truncated rows)', ds_tokens.num_rows)
+    assert len(ds_tokens), f'No examples left after filtering out truncated rows, try a longer max_length than {max_length}'
     return ds_tokens

src/repe/rep_control_pipeline_baukit.py

@@ -71,19 +71,35 @@ class RepControlPipeline2(FeatureExtractionPipeline):
 
     def __call__(self, model_inputs, activations=None, **kwargs):
         with torch.no_grad():
-            if activations is not None:
-                activations_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()})
+            if activations is not None:                
                 layers_names = [self.layer_name_tmpl.format(i) for i in activations.keys()]
-                edit_fn = partial(intervention_meta_fn2, activations=activations_i)
+                
+                # make intervention functions
+                activations_pos_i = Activations({self.layer_name_tmpl.format(k):v for k,v in activations.items()})
+                activations_neg_i = Activations({self.layer_name_tmpl.format(k):-v for k,v in activations.items()})
+                
+                edit_fn = partial(intervention_meta_fn2, activations=activations_pos_i)
                 with TraceDict(
                     self.model, layers_names, detach=True, edit_output=edit_fn
                 ) as ret:
-                    outputs = super().__call__(model_inputs, **kwargs)
+                    outputs_pos = super().__call__(model_inputs, **kwargs)
+                    
+                edit_fn2 = partial(intervention_meta_fn2, activations=activations_neg_i)
+                with TraceDict(
+                    self.model, layers_names, detach=True, edit_output=edit_fn2
+                ) as ret:
+                    outputs_neg = super().__call__(model_inputs, **kwargs)
+                    
+                outputs = super().__call__(model_inputs, **kwargs)
+                
+                # TODO stack the hidden states, and scores
+                pass
+                
             else:
                 outputs = super().__call__(model_inputs, **kwargs)
         return outputs
     
-    def preprocess(self, inputs: Dataset, **tokenize_kwargs) -> Dict[str, GenericTensor]:
+    def preprocess(self, inputs: dict, **tokenize_kwargs) -> Dict[str, GenericTensor]:
         # tokenize a batch of inputs
         return_tensors = self.framework
         
@@ -101,12 +117,12 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         inputs["attention_mask"] = torch.tensor(inputs['attention_mask'], dtype=torch.bool, device=self.model.device)
         return inputs
 
-    def _forward(self, inputs):
+    def _forward(self, inputs) -> ModelOutput:
         
         assert inputs['input_ids'].ndim == 2, f"expected input_ids to be (batch, seq), got {inputs['input_ids'].shape}"
         
         self.model.eval()
-        model_inputs = dict(
+        model_in = dict(
             input_ids=inputs['input_ids'],
             attention_mask=inputs['attention_mask'],
             use_cache=False,
@@ -114,18 +130,16 @@ class RepControlPipeline2(FeatureExtractionPipeline):
             return_dict=True
         )
         with torch.no_grad():
-            model_outputs = self.model(**model_inputs)
+            o = self.model(**model_in)
         
         # hidden states come at as lists of layers, lets concat them
-        model_outputs['hidden_states'] = rearrange(list(model_outputs['hidden_states']), 'l b t h -> b l t h')
-        model_outputs['last_hidden_states'] = model_outputs['hidden_states'][:, -1]
+        o['hidden_states'] = rearrange(list(o['hidden_states']), 'l b t h -> b l t h')
         
         # batch of outputs and inputs. retain some of the inputs
-        model_outputs = hacky_sanitize_outputs(model_outputs)
-        inputs = hacky_sanitize_outputs(inputs)
-        return ModelOutput(**model_outputs, **inputs)
+        return ModelOutput(**o, **inputs)
 
-    def postprocess(self, o):
+    def postprocess(self, o: ModelOutput):
+        o = hacky_sanitize_outputs(o)
         # note this sometimes deals with a batch, sometimes with a single result. infuriating
         res = []
         for i in range(len(o['input_ids'])):
@@ -138,10 +152,11 @@ class RepControlPipeline2(FeatureExtractionPipeline):
         else:
             return res
         
-    def postprocess_single(self, o):
+    def postprocess_single(self, o: dict) -> dict:
         assert isinstance(o, dict) and o['logits'].ndim==2, f"expected dict with logits of shape (seq, vocab), got {o['logits'].shape}"
         # assert o['logits'].shape[0]==1, f"postprocess expected batch size 1, got {o['logits'].shape[0]}"
         # This is called once for each result, but the text pipeline is set up to hande multiple...
+        o['last_hidden_states'] = o['hidden_states'][:, -1]
         
         o["end_logits"] = o["logits"][-1, :].float()