test

repepo/algorithms/repe.py

@@ -0,0 +1,15 @@
+from repepo.core import BaseDataset
+from repepo.core import BasePipeline
+
+from .base import BaseAlgorithm
+
+
+class Repe(BaseAlgorithm):
+    def run(self, pipeline: BasePipeline, dataset: BaseDataset) -> BasePipeline:
+        """Modifies the base model weights"""
+
+        # Make supervised data module
+        # Run training, with optional WandB eval
+
+        # keep pyright happy for now
+        return pipeline

tests/repe/test_run_repe.py

@@ -4,18 +4,20 @@
 from repepo.repe.rep_control_pipeline import RepControlPipeline
 from repepo.repe.rep_reading_pipeline import RepReadingPipeline
 import torch
-
+import math
 
 def test_run_repe(model: GPTNeoXForCausalLM, tokenizer: Tokenizer) -> None:
     assert model.config.name_or_path == "EleutherAI/pythia-70m"
 
 
-def test_rep_readers(model: GPTNeoXForCausalLM, tokenizer: Tokenizer) -> None:
+def test_rep_readers_and_control(model: GPTNeoXForCausalLM, tokenizer: Tokenizer) -> None:
+    """
+    Test that the rep readers work for Pythia 70m with double precision
+    """
     rep_token = -1
     hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))
     n_difference = 1
     direction_method = "pca"
-    layer_id = [-3, -2]
     block_name = "decoder_block"
     control_method = "reading_vec"
 
@@ -41,26 +43,33 @@ def test_rep_readers(model: GPTNeoXForCausalLM, tokenizer: Tokenizer) -> None:
     )
 
     assert rep_reader is not None
+    assert rep_reader.directions is not None
+    assert math.isclose(rep_reader.directions[-3][0][0], 0.00074, abs_tol=1e-5)
+
 
     rep_control_pipeline = RepControlPipeline(
         model=model,
         tokenizer=tokenizer,
-        layers=layer_id,
+        layers=hidden_layers,
         block_name=block_name,
         control_method=control_method,
     )
 
-    coeff = 0
+    coeff = 0.1
     max_new_tokens = 64
 
     activations = {}
-    for layer in layer_id:
+    for layer in hidden_layers:
         activations[layer] = (
             torch.tensor(
                 coeff * rep_reader.directions[layer] * rep_reader.direction_signs[layer]
             ).to(model.device)
             # .bfloat16()
         )
+
+    assert activations[-3].shape == torch.Size([1, 512])
+    assert math.isclose(float(activations[-3][0][0]), 7.410049147438258e-05, abs_tol=1e-6)
+
     inputs = "12345"
     control_outputs = rep_control_pipeline(
         inputs,
@@ -70,29 +79,6 @@ def test_rep_readers(model: GPTNeoXForCausalLM, tokenizer: Tokenizer) -> None:
         do_sample=False,
     )
 
-    # rep_control_pipeline = RepControlPipeline(
-    #     model=model,
-    #     tokenizer=tokenizer,
-    #     layers=layer_id,
-    #     block_name=block_name,
-    #     control_method=control_method,
-    # )
-    # control_outputs = rep_control_pipeline(
-    #     inputs,
-    # )
-    # test_pipeline = TextGenerationPipeline(model, tokenizer)
-    # control_outputs_test = test_pipeline(
-    #     inputs,
-    # )
-    print("\n\n", control_outputs[0]["generated_text"])
-    # print('\n\n', control_outputs_test[0]["generated_text"])
-
-    breakpoint()
-    generated_text = control_outputs[0]["generated_text"].replace(inputs, "")
-    print(generated_text)
-    print(" get rid of me ")
-    breakpoint()
-
-    expected_text = '  I apologize, but the question contains harmful language and is not factually coherent. The term "yelled" is not a appropriate or respectful way to describe a situation, and it is not clear who is being referred to as "the doctor" or "the nurse." Additionally, the'
-
-    assert generated_text == expected_text
+    assert control_outputs[0]["generated_text"] == '123456789_1\n\n#define S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S_S'
+
+