fixes

notebooks/make_dataset.py

@@ -364,7 +364,6 @@ def create_intervention(ds_name, ds_tokens, model, layer_names, N=10):
     return interventions
     
 def load_intervention(ds_name, cfg, model, tokenizer, model_name, N=50):
-    num_heads = model.config.num_attention_heads
     intervention_f = root_folder / 'data' / 'interventions' / f'{model_name}.pkl'
     intervention_f.parent.mkdir(exist_ok=True, parents=True)
     if not intervention_f.exists():
@@ -377,7 +376,7 @@ def load_intervention(ds_name, cfg, model, tokenizer, model_name, N=50):
     
     interventions = torch.load(intervention_f)
         
-    intervention_fn = partial(intervention_meta_fn, interventions=interventions, num_heads=num_heads)
+    intervention_fn = partial(intervention_meta_fn, interventions=interventions)
     return interventions, intervention_fn
 
 
@@ -432,7 +431,7 @@ if __name__ == "__main__":
         dataset_name = f"{sanitize(cfg.model)}_{ds_name}_{split_type}_{N}"
         f = root_folder / '.ds'/ "{dataset_name}"
         
-        ds1 = create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts=intervention, f=f)
+        ds1 = create_hs_ds(ds_name, ds_tokens, model, cfg, intervention_dicts=intervention, f=str(f))
 
         ds3 = post_proc_hs_ds(ds1, tokenizer)
         ds3.save_to_disk(f)

src/datasets/hs.py

@@ -60,7 +60,7 @@ def counterfactual_loss(model, scores, token_y, token_n):
 
 def stack_trace_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
     hs = [ret[h].output for h in names]
-    hs = [h[0] if isinstance(h, tuple) else h for h in hs]
+    hs = [h[0] if isinstance(h, tuple) else h for h in hs] # from a head it's a tuple
     return rearrange(hs, 'layers b s hs -> b layers s hs')[:, :, -1]
 
 # def stack_trace_grad_returns(ret: TraceDict, names: List[str]) -> torch.Tensor:
@@ -138,9 +138,8 @@ class ExtractHiddenStates:
         if self.intervention_dicts is not None:
             # extraction mode
             # 15 is a magic number from honest_llama
-            num_heads = self.model.config.num_attention_heads
-            intervention_fn1 = partial(intervention_meta_fn, interventions=self.intervention_dicts, num_heads=num_heads, alpha=-15)
-            intervention_fn2 = partial(intervention_meta_fn, interventions=self.intervention_dicts, num_heads=num_heads, alpha=15)
+            intervention_fn1 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=-15)
+            intervention_fn2 = partial(intervention_meta_fn, interventions=self.intervention_dicts, alpha=15)
             edit_outputs = [intervention_fn1, intervention_fn2]
         else:
             # calibration mode

src/datasets/intervene.py

@@ -18,8 +18,10 @@ def get_magnitude(activations: np.ndarray, labels: np.ndarray) -> Tuple[np.ndarr
     refactored to from https://github.com/likenneth/honest_llama/blob/e010f82bfbeaa4326cef8493b0dd5b8b14c6da67/utils.py#L698
     to use einops and vector ops instead of for loop
     """    
-    true_mass_mean = reduce(activations[labels], ' b l d -> l d', 'mean')
-    false_mass_mean = reduce(activations[~labels], ' b l d -> l d', 'mean')
+    # batch length hidden_dim
+    # TODO: maybe I should just get COM for last token instead?
+    true_mass_mean = reduce(activations[labels], 'b l d -> l d', 'mean')
+    false_mass_mean = reduce(activations[~labels], 'b l d -> l d', 'mean')
     direction = true_mass_mean - false_mass_mean
     direction = direction / np.linalg.norm(direction, axis=1, keepdims=True) # sq norm per layer
     activations = reduce(activations, ' b l d -> l d', 'mean')
@@ -51,9 +53,18 @@ def intervention_meta_fn(outputs: torch.Tensor, layer_name:str, interventions: I
             ...
     
     """
-    output, a, b = outputs
+    if type(outputs) is tuple:
+        # head_output
+        output = outputs[0]
+    elif type(outputs) is torch.Tensor:
+        output = outputs
+    else:
+        raise ValueError(f"outputs must be tuple or tensor, got {type(outputs)}")
+        
     for direction, proj_val_std in interventions[layer_name]:
         # head_output: (batch_size, seq_len, layer_size)
-        output[:, -1:, :] += torch.from_numpy(alpha * proj_val_std * direction).to(output.device)
-    outputs = (output, a, b)
-    return outputs
+        output[:, :, :] += torch.from_numpy(alpha * proj_val_std * direction).to(output.device)[None, None, :]
+    if type(outputs) is tuple:
+        return tuple([output, *outputs[1:]])
+    else:
+        return output