Generalization experiments

pyproject.toml

@@ -26,6 +26,13 @@ dependencies = [
     "nbdime>=4.0.1",
     "steering-vectors>=0.3.0",
     "openai>=1.10.0",
+    "arrr>=1.0.4",
+    "spacy>=3.7.2",
+    "mosestokenizer>=1.2.1",
+    "gradio>=4.16.0",
+    "simple-parsing>=0.1.5",
+    "torchmetrics>=1.3.0.post0",
+    "umap-learn>=0.5.5",
 ]
 
 [tool.black]

repepo/algorithms/repe.py

@@ -3,6 +3,7 @@
 from typing import Literal, Optional
 from typing_extensions import override
 import random
+import torch
 from repepo.core.pipeline import PipelineContext
 
 from steering_vectors import (
@@ -97,6 +98,7 @@ class RepeReadingControl(Algorithm):
     read_token_index: int
     seed: int
     show_progress: bool
+    verbose: bool
 
     def __init__(
         self,
@@ -117,6 +119,8 @@ def __init__(
         # Reference: https://github.com/nrimsky/SycophancySteering/blob/25f93a1f1aad51f94288f52d01f6a10d10f42bf1/generate_vectors.py#L102C13-L102C67
         read_token_index: int = -1,
         show_progress: bool = True,
+        verbose: bool = False,
+        steering_vector_save_path: Optional[str] = None,
     ):
         self.multi_answer_method = multi_answer_method
         self.layer_type = layer_type
@@ -130,6 +134,8 @@ def __init__(
         self.layer_config = layer_config
         self.direction_multiplier = direction_multiplier
         self.show_progress = show_progress
+        self.verbose = verbose
+        self.steering_vector_save_path = steering_vector_save_path
 
         self.skip_reading = skip_reading
         self.override_vector = override_vector
@@ -164,6 +170,17 @@ def _get_steering_vector(
         repe_training_data = self._build_steering_vector_training_data(
             dataset, pipeline
         )
+        if self.verbose:
+            # Print a small section of the dataset
+            pos_example, neg_example = repe_training_data[0]
+            print("Example steering vector training data:")
+            print("Positive prompt:")
+            print(pos_example)
+            print("Negative prompt:")
+            print(neg_example)
+            for i in range(2):
+                print()
+
         return train_steering_vector(
             pipeline.model,
             pipeline.tokenizer,
@@ -178,9 +195,7 @@ def _get_steering_vector(
 
     @override
     def run(self, pipeline: Pipeline, dataset: Dataset) -> Pipeline:
-        # Steering vector reading
-        # NOTE: The hooks read from this steering vector.
-
+        # TODO: Clean up this horrible code...
         if self.override_vector is not None:
             steering_vector: SteeringVector = self.override_vector
         elif not self.skip_reading:
@@ -192,6 +207,10 @@ def run(self, pipeline: Pipeline, dataset: Dataset) -> Pipeline:
                 "Either reading or override vector must be provided for control"
             )
 
+        if self.steering_vector_save_path is not None:
+            # TODO: Refactor into steering_vector.save
+            torch.save(steering_vector, self.steering_vector_save_path)
+
         # Creating the hooks that will do steering vector control
         # NOTE: How this works is that we create a context manager that creates a hook
         # whenever we are in a `PipelineContext`'s scope.

repepo/core/benchmark.py

@@ -54,6 +54,7 @@ def evaluate_benchmark(
     eval_hooks: list[EvalHook] = [],
     show_progress: bool = True,
     tqdm_desc: str = "Evaluating",
+    verbose: bool = False,
 ) -> EvalResult:
     # evaluate
     return evaluate(
@@ -64,6 +65,7 @@ def evaluate_benchmark(
         eval_hooks=eval_hooks,
         show_progress=show_progress,
         tqdm_desc=tqdm_desc,
+        verbose=verbose,
     )
 
 

repepo/core/evaluate.py

@@ -19,6 +19,20 @@
 EvalHook = Callable[[Pipeline], AbstractContextManager[None]]
 
 
+def print_first_example() -> EvalHook:
+    """Eval hook that prints the first example"""
+
+    @contextmanager
+    def print_first_example_hook(pipeline: Pipeline):
+        try:
+            pipeline.print_first_example = True
+            yield
+        finally:
+            pipeline.print_first_example = False
+
+    return print_first_example_hook
+
+
 def update_completion_template_at_eval(new_template: str) -> EvalHook:
     """Eval hook that changes the completion template for the duration of the evaluation"""
 
@@ -200,7 +214,7 @@ def score_prediction(self, prediction: EvalPrediction) -> float:
 
     def __call__(self, predictions: Sequence[EvalPrediction]) -> dict[str, float]:
         pred_results = [self.score_prediction(pred) for pred in predictions]
-        return {"accuracy": mean(pred_results)}
+        return {"average_key_prob": mean(pred_results)}
 
 
 def evaluate(
@@ -213,6 +227,7 @@ def evaluate(
     eval_hooks: Sequence[EvalHook] = [],
     show_progress: bool = True,
     tqdm_desc: str = "Evaluating",
+    verbose: bool = False,
 ) -> EvalResult:
     # evaluate
     predictions: list[EvalPrediction] = []
@@ -222,15 +237,23 @@ def evaluate(
         for eval_hook in eval_hooks:
             stack.enter_context(eval_hook(pipeline))
         # TODO: support batching
-        for example in tqdm(dataset, disable=not show_progress, desc=tqdm_desc):
+        for i, example in enumerate(
+            tqdm(dataset, disable=not show_progress, desc=tqdm_desc)
+        ):
             generated_output = None
             correct_output_probs = None
             incorrect_outputs_probs = None
             if requires_generation:
+                if i == 0 and verbose:
+                    print("Example generation prompt:")
+                    print(pipeline.build_generation_prompt(example))
                 generated_output = pipeline.generate(
                     example, generation_config=generation_config
                 )
             if requires_probs:
+                if i == 0 and verbose:
+                    print("Example full prompt:")
+                    print(pipeline.build_full_prompt(example))
                 correct_output_probs = pipeline.calculate_output_logprobs(example)
                 if example.incorrect_outputs is not None:
                     incorrect_outputs_probs = [

repepo/core/pipeline.py

@@ -50,6 +50,7 @@ class Pipeline:
     formatter: Formatter = field(default_factory=InputOutputFormatter)
     conversation_history: list[Example] = field(default_factory=list)
     hooks: list[PipelineHook] = field(default_factory=list)
+    print_first_example: bool = True
 
     def build_generation_prompt(self, example: Example) -> str:
         """Build a prompt for generation"""

repepo/data/make_dataset.py

@@ -23,15 +23,17 @@ def get_all_json_filepaths(root_dir: pathlib.Path) -> List[pathlib.Path]:
 
 
 # Intentionally don't cache anything here, otherwise datasets don't be available after downloading
-def _get_datasets() -> dict[str, pathlib.Path]:
+def _get_datasets(dataset_dir: pathlib.Path | None = None) -> dict[str, pathlib.Path]:
     datasets: dict[str, pathlib.Path] = {}
-    for path in get_all_json_filepaths(get_dataset_dir()):
+    if dataset_dir is None:
+        dataset_dir = get_dataset_dir()
+    for path in get_all_json_filepaths(dataset_dir):
         datasets[path.stem] = path.absolute()
     return datasets
 
 
-def list_datasets() -> tuple[str, ...]:
-    return tuple(_get_datasets().keys())
+def list_datasets(dataset_dir: pathlib.Path | None = None) -> tuple[str, ...]:
+    return tuple(_get_datasets(dataset_dir).keys())
 
 
 @dataclass
@@ -40,6 +42,9 @@ class DatasetSpec:
     split: str = ":100%"
     seed: int = 0
 
+    def __repr__(self) -> str:
+        return f"DatasetSpec(name={self.name},split={self.split},seed={self.seed})"
+
 
 def _parse_split(split_string: str, length: int) -> slice:
     # Define the regular expression pattern
@@ -61,8 +66,8 @@ def _parse_split(split_string: str, length: int) -> slice:
         raise ValueError(f"Parse string {split_string} not recognized")
 
 
-def get_dataset(name: str) -> Dataset:
-    datasets = _get_datasets()
+def get_dataset(name: str, dataset_dir: pathlib.Path | None = None) -> Dataset:
+    datasets = _get_datasets(dataset_dir)
     if name not in datasets:
         raise ValueError(f"Unknown dataset: {name}")
 
@@ -84,6 +89,6 @@ def _shuffle_and_split(items: list[T], split_string: str, seed: int) -> list[T]:
     return shuffled_items[split]
 
 
-def make_dataset(spec: DatasetSpec):
-    dataset = get_dataset(spec.name)
+def make_dataset(spec: DatasetSpec, dataset_dir: pathlib.Path | None = None):
+    dataset = get_dataset(spec.name, dataset_dir)
     return _shuffle_and_split(dataset, spec.split, spec.seed)