add token score results

Author: mortalhermit

linear-probes/all_layers_cache_train_eval.py

@@ -27,7 +27,7 @@
 import probes
 from pprint import pprint as pp
 
-from datasets import AmongUsDataset, TruthfulQADataset, DishonestQADataset, RepEngDataset, RolePlayingDataset, ApolloProbeDataset
+from probe_datasets import AmongUsDataset, TruthfulQADataset, DishonestQADataset, RepEngDataset, RolePlayingDataset, ApolloProbeDataset
 from configs import config_phi4, config_gpt2, config_llama3
 base_config = config_phi4
 amongus_expt_name: str = "2025-02-01_phi_phi_100_games_v3"

linear-probes/cache_activations.py

@@ -11,7 +11,7 @@
 for module in [datasets, plots, configs, evaluate_utils]:
     importlib.reload(module)
 
-from datasets import AmongUsDataset, TruthfulQADataset, DishonestQADataset, RepEngDataset, RolePlayingDataset, ApolloProbeDataset
+from probe_datasets import AmongUsDataset, TruthfulQADataset, DishonestQADataset, RepEngDataset, RolePlayingDataset, ApolloProbeDataset
 from configs import config_phi4, config_gpt2, config_llama3
 
 def main(dataset_name: str):

linear-probes/evaluate_probes.py

@@ -15,7 +15,7 @@
 from typing import Dict, Any, List, Tuple
 from sklearn.metrics import roc_curve, auc, precision_recall_curve, average_precision_score
 
-from datasets import TruthfulQADataset, DishonestQADataset, AmongUsDataset, RolePlayingDataset, RepEngDataset
+from probe_datasets import TruthfulQADataset, DishonestQADataset, AmongUsDataset, RolePlayingDataset, RepEngDataset
 from evaluate_utils import evaluate_probe_on_activation_dataset
 from configs import config_phi4, config_gpt2, config_llama3
 from plots import plot_behavior_distribution, plot_roc_curves, add_roc_curves, print_metrics, plot_roc_curve_eval

linear-probes/evaluate_utils.py

@@ -75,7 +75,7 @@ def evaluate_probe_on_dataset(test_df, model, tokenizer, probe, dataset, device,
     print(f"Accuracy: {accuracy}")
     return av_probe_outputs, accuracy
 
-def evaluate_probe_on_activation_dataset(chunk_data, probe, device, num_tokens=None, verbose=True):
+def evaluate_probe_on_activation_dataset(chunk_data, probe, device, num_tokens=None, verbose=True, output_token_scores=False):
     """
     Evaluate probe on a test dataset of activations and return probe outputs and accuracy
     
@@ -88,34 +88,41 @@ def evaluate_probe_on_activation_dataset(chunk_data, probe, device, num_tokens=N
         verbose: Whether to print progress
         
     Returns:
-        tuple: (average probe outputs, accuracy)
+        tuple: (average probe outputs, accuracy, probe outputs) if output_token_scores is True, otherwise (average probe outputs, accuracy)
     """
     av_probe_outputs = []
+    all_probe_outputs = []
     total, correct = 0, 0
 
-    for i, (activations, label) in enumerate(chunk_data):
+    from tqdm import tqdm
+    
+    for i, (activations, label) in enumerate(tqdm(chunk_data)):
         total += 1
 
         # Take specified number of tokens from end of sequence
         acts_to_use = activations[-num_tokens:] if num_tokens else activations
 
         # Get probe output for each token position
-        probe_outputs = [round(probe.evaluate_single_activation(t.tensor(act, device=device)), 4) 
+        token_probe_outputs = [round(probe.evaluate_single_activation(t.tensor(act, device=device)), 4) 
                         for act in acts_to_use]
-        avg_probe_output = sum(probe_outputs) / len(probe_outputs)
+        avg_probe_output = sum(token_probe_outputs) / len(token_probe_outputs)
 
         if label == 1 and avg_probe_output > 0.5:
             correct += 1
         elif label == 0 and avg_probe_output <= 0.5:
             correct += 1
 
         av_probe_outputs.append(avg_probe_output)
+        all_probe_outputs.append(token_probe_outputs)
 
         if verbose and i % ((len(chunk_data) + 10) // 10) == 0:
             print(f"Evaluating {i}/{len(chunk_data)}", end="\t")
-            print(f"Probe outputs: {probe_outputs}")
+            print(f"Probe outputs: {token_probe_outputs}")
 
     accuracy = correct / total
     if verbose:
         print(f".", end="")
-    return av_probe_outputs, accuracy
+    if output_token_scores:
+        return av_probe_outputs, accuracy, all_probe_outputs
+    else:
+        return av_probe_outputs, accuracy

linear-probes/datasets.py linear-probes/probe_datasets.pylinear-probes/datasets.py renamed to linear-probes/probe_datasets.py

@@ -72,10 +72,11 @@ def __init__(self, test_split, name: str = "", model=None, tokenizer=None, devic
             self.activation_cache.remove_hooks()
 
     def get_chunk_path(self, chunk_idx: int) -> str:
-        return os.path.join(self.activations_dir, f"chunk_{chunk_idx}.pkl")
+        return os.path.join(os.path.dirname(os.path.abspath(__file__)), self.activations_dir, f"chunk_{chunk_idx}.pkl")
 
     def save_chunk(self, chunk_data: List[Tuple[List[t.Tensor], int]], chunk_idx: int):
-        os.makedirs(self.activations_dir, exist_ok=True)
+        full_dir_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.activations_dir)
+        os.makedirs(full_dir_path, exist_ok=True)
         chunk_path = self.get_chunk_path(chunk_idx)
         with open(chunk_path, 'wb') as f:
             pickle.dump(chunk_data, f)
@@ -155,9 +156,9 @@ def get_train_data_stats(self, chunk_idx: int = 0) -> dict:
 class TruthfulQADataset(ActivationDataset):
     def __init__(self, config: Dict[str, Any]=None, model=None, tokenizer=None, device=None, test_split=None, **kwargs):
         super().__init__(test_split, "TruthfulQA", model, tokenizer, device, config["activation_size"])
-        self.data_path: str = './data/TruthfulQA/TruthfulQA.csv'
+        self.data_path: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), './data/TruthfulQA/TruthfulQA.csv')
         self.tqa_df = pd.read_csv(self.data_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         self.num_total_chunks = 1  # TruthfulQA uses single chunk
         self.format = eval(config["short_name"] + "_format")
 
@@ -203,9 +204,9 @@ def populate_dataset(self, force_redo: bool = False, num_tokens: int = 5, max_ro
 class DishonestQADataset(ActivationDataset):
     def __init__(self, config: Dict[str, Any]=None, model=None, tokenizer=None, device=None, test_split=None, **kwargs):
         super().__init__(test_split, "DishonestQA", model, tokenizer, device, config["activation_size"])
-        self.data_path: str = './data/TruthfulQA/TruthfulQA.csv'
+        self.data_path: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), './data/TruthfulQA/TruthfulQA.csv')
         self.tqa_df = pd.read_csv(self.data_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         self.num_total_chunks = 1  # DishonestQA uses single chunk
         self.format = eval(config["short_name"] + "_format")
 
@@ -272,11 +273,12 @@ def __init__(
             ):
         super().__init__(test_split, "AmongUs", model, tokenizer, device, config["activation_size"])
         self.name: str = "AmongUs"
-        self.agent_logs_path: str = os.path.join(raw_path, expt_name + "/agent-logs-compact.json")
-        sys.path.append("..")
+        base_dir = os.path.dirname(os.path.abspath(__file__))
+        self.agent_logs_path: str = os.path.join(base_dir, raw_path, expt_name + "/agent-logs-compact.json")
+        sys.path.append(os.path.join(base_dir, ".."))
         from utils import load_agent_logs_df
         self.agent_logs_df = load_agent_logs_df(self.agent_logs_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         # load number of chunks from existing directory
         self.num_total_chunks = 0
         self.format = eval(config["short_name"] + "_format")
@@ -335,11 +337,14 @@ def populate_dataset(
             print(f"Loaded {self.num_total_chunks} existing chunks")
             return
 
-        if force_redo and os.path.exists(self.activations_dir):
-            import shutil
-            shutil.rmtree(self.activations_dir)
+        if force_redo:
+            full_dir_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.activations_dir)
+            if os.path.exists(full_dir_path):
+                import shutil
+                shutil.rmtree(full_dir_path)
 
-        os.makedirs(self.activations_dir, exist_ok=True)
+        full_dir_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), self.activations_dir)
+        os.makedirs(full_dir_path, exist_ok=True)
 
         # Find last processed chunk
         chunk_idx = 0
@@ -383,9 +388,9 @@ def populate_dataset(
 class RolePlayingDataset(ActivationDataset):
     def __init__(self, config: Dict[str, Any]=None, model=None, tokenizer=None, device=None, test_split=None, **kwargs):
         super().__init__(test_split, "Roleplaying", model, tokenizer, device, config["activation_size"])
-        self.data_path: str = './data/Roleplaying/phi4_rollouts.csv'
+        self.data_path: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), './data/Roleplaying/phi4_rollouts.csv')
         self.df = pd.read_csv(self.data_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         self.num_total_chunks = 1  # Roleplaying uses single chunk
         self.format = eval(config["short_name"] + "_format")
 
@@ -424,9 +429,9 @@ def populate_dataset(self, force_redo: bool = False, num_tokens: int = 5, max_ro
 class RepEngDataset(ActivationDataset):
     def __init__(self, config: Dict[str, Any]=None, model=None, tokenizer=None, device=None, test_split=None, **kwargs):
         super().__init__(test_split, "RepEng", model, tokenizer, device, config["activation_size"])
-        self.data_path: str = './data/RepE/true_false_facts.csv'
+        self.data_path: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), './data/RepE/true_false_facts.csv')
         self.df = pd.read_csv(self.data_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         self.num_total_chunks = 1 # RepEng uses single chunk
         self.format = eval(config["short_name"] + "_format")
 
@@ -465,9 +470,9 @@ def populate_dataset(self, force_redo: bool = False, num_tokens: int = 5, max_ro
 class ApolloProbeDataset(ActivationDataset):
     def __init__(self, config: Dict[str, Any]=None, model=None, tokenizer=None, device=None, test_split=None, **kwargs):
         super().__init__(test_split, "ApolloProbe", model, tokenizer, device, config["activation_size"])
-        self.data_path: str = './data/ApolloProbe/common_claim_true_false.csv'
+        self.data_path: str = os.path.join(os.path.dirname(os.path.abspath(__file__)), './data/ApolloProbe/common_claim_true_false.csv')
         self.df = pd.read_csv(self.data_path)
-        self.activations_dir: str = f'./data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
+        self.activations_dir: str = f'data/{self.name}_{config["short_name"]}_acts_{config["layer"]}/'
         self.num_total_chunks = 1 # ApolloProbe uses single chunk
         self.format = eval(config["short_name"] + "_format")
 

linear-probes/results/RepEngDataset_phi4/layer_20_roc_AmongUs.pdf

@@ -12,7 +12,7 @@
 sys.path.append(os.path.dirname(os.path.abspath('.')))
 sys.path.append('.')
 
-from datasets import (
+from probe_datasets import (
     TruthfulQADataset,
     DishonestQADataset, 
     AmongUsDataset,

linear-probes/train_probes.py

@@ -4285,7 +4285,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [