clean up result analysis

Author: luciaquirke

delphi/__main__.py

@@ -55,7 +55,12 @@ def load_artifacts(run_cfg: RunConfig):
         compile=True,
     )
 
-    return run_cfg.hookpoints, hookpoint_to_sparse_encode, model, transcode
+    return (
+        list(hookpoint_to_sparse_encode.keys()),
+        hookpoint_to_sparse_encode,
+        model,
+        transcode,
+    )
 
 
 def create_neighbours(
@@ -212,7 +217,7 @@ def scorer_postprocess(result, score_dir):
                 client,
                 n_examples_shown=run_cfg.num_examples_per_scorer_prompt,
                 verbose=run_cfg.verbose,
-                log_prob=False,
+                log_prob=run_cfg.log_probs,
             ),
             preprocess=scorer_preprocess,
             postprocess=partial(scorer_postprocess, score_dir=detection_scores_path),
@@ -222,7 +227,7 @@ def scorer_postprocess(result, score_dir):
                 client,
                 n_examples_shown=run_cfg.num_examples_per_scorer_prompt,
                 verbose=run_cfg.verbose,
-                log_prob=False,
+                log_prob=run_cfg.log_probs,
             ),
             preprocess=scorer_preprocess,
             postprocess=partial(scorer_postprocess, score_dir=fuzz_scores_path),

delphi/config.py

@@ -142,6 +142,9 @@ class RunConfig(Serializable):
     filter_bos: bool = False
     """Whether to filter out BOS tokens from the cache."""
 
+    log_probs: bool = False
+    """Whether to attempt to gather log probabilities for each scorer prompt."""
+
     load_in_8bit: bool = False
     """Load the model in 8-bit mode."""